Merge branch 'main' of https://github.com/danswer-ai/danswer into feature/background_status

* Mismatch issue of Documents shown and Citation number in text fix

When document order presented to LLM differs from order shown to user, wrong doc numbers are cited.

Fix:
 - SearchTool.get_search_result  returns now final and initial ranking
 - initial ranking is passed through a few objects and used for replacement in citation processing

Notes:
 - the citation_num in the CitationInfo() object has not been changed.

* PR fixes

 - linting
 - removed erroneous tab
 - added a substitution test case
 - adjusted original citation extraction use case

* Included a key test and

* Fixed extra spaces

* Updated test documentation

Updated:
 - test_citation_substitution (changed description)
 - test_citation_processing (removed data only relevant for the substitution)

README.md

@@ -1,48 +1,48 @@
-<!-- DANSWER_METADATA={"link": "https://github.com/danswer-ai/danswer/blob/main/README.md"} -->
+<!-- DANSWER_METADATA={"link": "https://github.com/onyx-dot-app/onyx/blob/main/README.md"} -->
 <a name="readme-top"></a>
 
 <h2 align="center">
-<a href="https://www.danswer.ai/"> <img width="50%" src="https://github.com/danswer-owners/danswer/blob/1fabd9372d66cd54238847197c33f091a724803b/DanswerWithName.png?raw=true)" /></a>
+<a href="https://www.onyx.app/"> <img width="50%" src="https://github.com/onyx-dot-app/onyx/blob/logo/LogoOnyx.png?raw=true)" /></a>
 </h2>
 
 <p align="center">
-<p align="center">Open Source Gen-AI Chat + Unified Search.</p>
+<p align="center">Open Source Gen-AI + Enterprise Search.</p>
 
 <p align="center">
-<a href="https://docs.danswer.dev/" target="_blank">
+<a href="https://docs.onyx.app/" target="_blank">
     <img src="https://img.shields.io/badge/docs-view-blue" alt="Documentation">
 </a>
-<a href="https://join.slack.com/t/danswer/shared_invite/zt-2twesxdr6-5iQitKZQpgq~hYIZ~dv3KA" target="_blank">
+<a href="https://join.slack.com/t/onyx-dot-app/shared_invite/zt-2sslpdbyq-iIbTaNIVPBw_i_4vrujLYQ" target="_blank">
     <img src="https://img.shields.io/badge/slack-join-blue.svg?logo=slack" alt="Slack">
 </a>
 <a href="https://discord.gg/TDJ59cGV2X" target="_blank">
     <img src="https://img.shields.io/badge/discord-join-blue.svg?logo=discord&logoColor=white" alt="Discord">
 </a>
-<a href="https://github.com/danswer-ai/danswer/blob/main/README.md" target="_blank">
+<a href="https://github.com/onyx-dot-app/onyx/blob/main/README.md" target="_blank">
     <img src="https://img.shields.io/static/v1?label=license&message=MIT&color=blue" alt="License">
 </a>
 </p>
 
-<strong>[Danswer](https://www.danswer.ai/)</strong> is the AI Assistant connected to your company's docs, apps, and people. 
-Danswer provides a Chat interface and plugs into any LLM of your choice. Danswer can be deployed anywhere and for any 
+<strong>[Onyx](https://www.onyx.app/)</strong> (Formerly Danswer) is the AI Assistant connected to your company's docs, apps, and people. 
+Onyx provides a Chat interface and plugs into any LLM of your choice. Onyx can be deployed anywhere and for any 
 scale - on a laptop, on-premise, or to cloud. Since you own the deployment, your user data and chats are fully in your 
-own control. Danswer is MIT licensed and designed to be modular and easily extensible. The system also comes fully ready 
+own control. Onyx is dual Licensed with most of it under MIT license and designed to be modular and easily extensible. The system also comes fully ready 
 for production usage with user authentication, role management (admin/basic users), chat persistence, and a UI for 
-configuring Personas (AI Assistants) and their Prompts.
+configuring AI Assistants.
 
-Danswer also serves as a Unified Search across all common workplace tools such as Slack, Google Drive, Confluence, etc.
-By combining LLMs and team specific knowledge, Danswer becomes a subject matter expert for the team. Imagine ChatGPT if
+Onyx also serves as a Enterprise Search across all common workplace tools such as Slack, Google Drive, Confluence, etc.
+By combining LLMs and team specific knowledge, Onyx becomes a subject matter expert for the team. Imagine ChatGPT if
 it had access to your team's unique knowledge! It enables questions such as "A customer wants feature X, is this already
 supported?" or "Where's the pull request for feature Y?"
 
 <h3>Usage</h3>
 
-Danswer Web App:
+Onyx Web App:
 
 https://github.com/danswer-ai/danswer/assets/32520769/563be14c-9304-47b5-bf0a-9049c2b6f410
 
 
-Or, plug Danswer into your existing Slack workflows (more integrations to come 😁):
+Or, plug Onyx into your existing Slack workflows (more integrations to come 😁):
 
 https://github.com/danswer-ai/danswer/assets/25087905/3e19739b-d178-4371-9a38-011430bdec1b
 
@@ -52,16 +52,16 @@ For more details on the Admin UI to manage connectors and users, check out our
 
 ## Deployment
 
-Danswer can easily be run locally (even on a laptop) or deployed on a virtual machine with a single
-`docker compose` command. Checkout our [docs](https://docs.danswer.dev/quickstart) to learn more.
+Onyx can easily be run locally (even on a laptop) or deployed on a virtual machine with a single
+`docker compose` command. Checkout our [docs](https://docs.onyx.app/quickstart) to learn more.
 
-We also have built-in support for deployment on Kubernetes. Files for that can be found [here](https://github.com/danswer-ai/danswer/tree/main/deployment/kubernetes).
+We also have built-in support for deployment on Kubernetes. Files for that can be found [here](https://github.com/onyx-dot-app/onyx/tree/main/deployment/kubernetes).
 
 
 ## 💃 Main Features 
 * Chat UI with the ability to select documents to chat with.
 * Create custom AI Assistants with different prompts and backing knowledge sets.
-* Connect Danswer with LLM of your choice (self-host for a fully airgapped solution).
+* Connect Onyx with LLM of your choice (self-host for a fully airgapped solution).
 * Document Search + AI Answers for natural language queries.
 * Connectors to all common workplace tools like Google Drive, Confluence, Slack, etc.
 * Slack integration to get answers and search results directly in Slack.
@@ -75,12 +75,12 @@ We also have built-in support for deployment on Kubernetes. Files for that can b
 * Organizational understanding and ability to locate and suggest experts from your team.
 
 
-## Other Notable Benefits of Danswer
+## Other Notable Benefits of Onyx
 * User Authentication with document level access management.
 * Best in class Hybrid Search across all sources (BM-25 + prefix aware embedding models).
 * Admin Dashboard to configure connectors, document-sets, access, etc.
 * Custom deep learning models + learn from user feedback.
-* Easy deployment and ability to host Danswer anywhere of your choosing.
+* Easy deployment and ability to host Onyx anywhere of your choosing.
 
 
 ## 🔌 Connectors
@@ -108,10 +108,10 @@ Efficiently pulls the latest changes from:
 
 ## 📚 Editions
 
-There are two editions of Danswer:
+There are two editions of Onyx:
 
-  * Danswer Community Edition (CE) is available freely under the MIT Expat license. This version has ALL the core features discussed above. This is the version of Danswer you will get if you follow the Deployment guide above.
-  * Danswer Enterprise Edition (EE) includes extra features that are primarily useful for larger organizations. Specifically, this includes:
+  * Onyx Community Edition (CE) is available freely under the MIT Expat license. This version has ALL the core features discussed above. This is the version of Onyx you will get if you follow the Deployment guide above.
+  * Onyx Enterprise Edition (EE) includes extra features that are primarily useful for larger organizations. Specifically, this includes:
     * Single Sign-On (SSO), with support for both SAML and OIDC
     * Role-based access control
     * Document permission inheritance from connected sources
@@ -119,24 +119,24 @@ There are two editions of Danswer:
     * Whitelabeling
     * API key authentication
     * Encryption of secrets
-    * Any many more! Checkout [our website](https://www.danswer.ai/) for the latest.
+    * Any many more! Checkout [our website](https://www.onyx.app/) for the latest.
 
-To try the Danswer Enterprise Edition: 
+To try the Onyx Enterprise Edition: 
 
-  1. Checkout our [Cloud product](https://app.danswer.ai/signup).
-  2. For self-hosting, contact us at [founders@danswer.ai](mailto:founders@danswer.ai) or book a call with us on our [Cal](https://cal.com/team/danswer/founders).
+  1. Checkout our [Cloud product](https://cloud.onyx.app/signup).
+  2. For self-hosting, contact us at [founders@onyx.app](mailto:founders@onyx.app) or book a call with us on our [Cal](https://cal.com/team/danswer/founders).
 
 ## 💡 Contributing
 Looking to contribute? Please check out the [Contribution Guide](CONTRIBUTING.md) for more details.
 
 ## ⭐Star History
 
-[![Star History Chart](https://api.star-history.com/svg?repos=danswer-ai/danswer&type=Date)](https://star-history.com/#danswer-ai/danswer&Date)
+[![Star History Chart](https://api.star-history.com/svg?repos=onyx-dot-app/onyx&type=Date)](https://star-history.com/#onyx-dot-app/onyx&Date)
 
 ## ✨Contributors
 
-<a href="https://github.com/danswer-ai/danswer/graphs/contributors">
-  <img alt="contributors" src="https://contrib.rocks/image?repo=danswer-ai/danswer"/>
+<a href="https://github.com/onyx-dot-app/onyx/graphs/contributors">
+  <img alt="contributors" src="https://contrib.rocks/image?repo=onyx-dot-app/onyx"/>
 </a>
 
 <p align="right" style="font-size: 14px; color: #555; margin-top: 20px;">

backend/alembic/versions/bf7a81109301_delete_input_prompts.py

@@ -0,0 +1,57 @@
+"""delete_input_prompts
+
+Revision ID: bf7a81109301
+Revises: f7a894b06d02
+Create Date: 2024-12-09 12:00:49.884228
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import fastapi_users_db_sqlalchemy
+
+
+# revision identifiers, used by Alembic.
+revision = "bf7a81109301"
+down_revision = "f7a894b06d02"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.drop_table("inputprompt__user")
+    op.drop_table("inputprompt")
+
+
+def downgrade() -> None:
+    op.create_table(
+        "inputprompt",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("prompt", sa.String(), nullable=False),
+        sa.Column("content", sa.String(), nullable=False),
+        sa.Column("active", sa.Boolean(), nullable=False),
+        sa.Column("is_public", sa.Boolean(), nullable=False),
+        sa.Column(
+            "user_id",
+            fastapi_users_db_sqlalchemy.generics.GUID(),
+            nullable=True,
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["user.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_table(
+        "inputprompt__user",
+        sa.Column("input_prompt_id", sa.Integer(), nullable=False),
+        sa.Column("user_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["input_prompt_id"],
+            ["inputprompt.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["inputprompt.id"],
+        ),
+        sa.PrimaryKeyConstraint("input_prompt_id", "user_id"),
+    )

backend/alembic/versions/f7a894b06d02_non_nullbale_slack_bot_id_in_channel_.py

@@ -0,0 +1,40 @@
+"""non-nullbale slack bot id in channel config
+
+Revision ID: f7a894b06d02
+Revises: 9f696734098f
+Create Date: 2024-12-06 12:55:42.845723
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "f7a894b06d02"
+down_revision = "9f696734098f"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Delete all rows with null slack_bot_id
+    op.execute("DELETE FROM slack_channel_config WHERE slack_bot_id IS NULL")
+
+    # Make slack_bot_id non-nullable
+    op.alter_column(
+        "slack_channel_config",
+        "slack_bot_id",
+        existing_type=sa.Integer(),
+        nullable=False,
+    )
+
+
+def downgrade() -> None:
+    # Make slack_bot_id nullable again
+    op.alter_column(
+        "slack_channel_config",
+        "slack_bot_id",
+        existing_type=sa.Integer(),
+        nullable=True,
+    )

backend/danswer/agent_search/answer_query/graph_builder.py

@@ -1,100 +0,0 @@
-from langgraph.graph import END
-from langgraph.graph import START
-from langgraph.graph import StateGraph
-
-from danswer.agent_search.answer_query.nodes.answer_check import answer_check
-from danswer.agent_search.answer_query.nodes.answer_generation import answer_generation
-from danswer.agent_search.answer_query.nodes.format_answer import format_answer
-from danswer.agent_search.answer_query.states import AnswerQueryInput
-from danswer.agent_search.answer_query.states import AnswerQueryOutput
-from danswer.agent_search.answer_query.states import AnswerQueryState
-from danswer.agent_search.expanded_retrieval.graph_builder import (
-    expanded_retrieval_graph_builder,
-)
-
-
-def answer_query_graph_builder() -> StateGraph:
-    graph = StateGraph(
-        state_schema=AnswerQueryState,
-        input=AnswerQueryInput,
-        output=AnswerQueryOutput,
-    )
-
-    ### Add nodes ###
-
-    expanded_retrieval = expanded_retrieval_graph_builder().compile()
-    graph.add_node(
-        node="expanded_retrieval_for_initial_decomp",
-        action=expanded_retrieval,
-    )
-    graph.add_node(
-        node="answer_check",
-        action=answer_check,
-    )
-    graph.add_node(
-        node="answer_generation",
-        action=answer_generation,
-    )
-    graph.add_node(
-        node="format_answer",
-        action=format_answer,
-    )
-
-    ### Add edges ###
-
-    graph.add_edge(
-        start_key=START,
-        end_key="expanded_retrieval_for_initial_decomp",
-    )
-    graph.add_edge(
-        start_key="expanded_retrieval_for_initial_decomp",
-        end_key="answer_generation",
-    )
-    graph.add_edge(
-        start_key="answer_generation",
-        end_key="answer_check",
-    )
-    graph.add_edge(
-        start_key="answer_check",
-        end_key="format_answer",
-    )
-    graph.add_edge(
-        start_key="format_answer",
-        end_key=END,
-    )
-
-    return graph
-
-
-if __name__ == "__main__":
-    from danswer.db.engine import get_session_context_manager
-    from danswer.llm.factory import get_default_llms
-    from danswer.context.search.models import SearchRequest
-
-    graph = answer_query_graph_builder()
-    compiled_graph = graph.compile()
-    primary_llm, fast_llm = get_default_llms()
-    search_request = SearchRequest(
-        query="Who made Excel and what other products did they make?",
-    )
-    with get_session_context_manager() as db_session:
-        inputs = AnswerQueryInput(
-            search_request=search_request,
-            primary_llm=primary_llm,
-            fast_llm=fast_llm,
-            db_session=db_session,
-            query_to_answer="Who made Excel?",
-        )
-        output = compiled_graph.invoke(
-            input=inputs,
-            # debug=True,
-            # subgraphs=True,
-        )
-        print(output)
-        # for namespace, chunk in compiled_graph.stream(
-        #     input=inputs,
-        #     # debug=True,
-        #     subgraphs=True,
-        # ):
-        #     print(namespace)
-        #     print(chunk)

backend/danswer/agent_search/answer_query/nodes/answer_check.py

@@ -1,30 +0,0 @@
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.answer_query.states import AnswerQueryState
-from danswer.agent_search.answer_query.states import QACheckOutput
-from danswer.agent_search.shared_graph_utils.prompts import BASE_CHECK_PROMPT
-
-
-def answer_check(state: AnswerQueryState) -> QACheckOutput:
-    msg = [
-        HumanMessage(
-            content=BASE_CHECK_PROMPT.format(
-                question=state["search_request"].query,
-                base_answer=state["answer"],
-            )
-        )
-    ]
-
-    fast_llm = state["fast_llm"]
-    response = list(
-        fast_llm.stream(
-            prompt=msg,
-        )
-    )
-
-    response_str = merge_message_runs(response, chunk_separator="")[0].content
-
-    return QACheckOutput(
-        answer_quality=response_str,
-    )

backend/danswer/agent_search/answer_query/nodes/answer_generation.py

@@ -1,32 +0,0 @@
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.answer_query.states import AnswerQueryState
-from danswer.agent_search.answer_query.states import QAGenerationOutput
-from danswer.agent_search.shared_graph_utils.prompts import BASE_RAG_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-
-
-def answer_generation(state: AnswerQueryState) -> QAGenerationOutput:
-    query = state["query_to_answer"]
-    docs = state["reordered_documents"]
-
-    print(f"Number of verified retrieval docs: {len(docs)}")
-
-    msg = [
-        HumanMessage(
-            content=BASE_RAG_PROMPT.format(question=query, context=format_docs(docs))
-        )
-    ]
-
-    fast_llm = state["fast_llm"]
-    response = list(
-        fast_llm.stream(
-            prompt=msg,
-        )
-    )
-
-    answer_str = merge_message_runs(response, chunk_separator="")[0].content
-    return QAGenerationOutput(
-        answer=answer_str,
-    )

backend/danswer/agent_search/answer_query/nodes/format_answer.py

@@ -1,16 +0,0 @@
-from danswer.agent_search.answer_query.states import AnswerQueryOutput
-from danswer.agent_search.answer_query.states import AnswerQueryState
-from danswer.agent_search.answer_query.states import SearchAnswerResults
-
-
-def format_answer(state: AnswerQueryState) -> AnswerQueryOutput:
-    return AnswerQueryOutput(
-        decomp_answer_results=[
-            SearchAnswerResults(
-                query=state["query_to_answer"],
-                quality=state["answer_quality"],
-                answer=state["answer"],
-                documents=state["reordered_documents"],
-            )
-        ],
-    )

backend/danswer/agent_search/answer_query/states.py

@@ -1,45 +0,0 @@
-from typing import Annotated
-from typing import TypedDict
-
-from pydantic import BaseModel
-
-from danswer.agent_search.core_state import PrimaryState
-from danswer.agent_search.shared_graph_utils.operators import dedup_inference_sections
-from danswer.context.search.models import InferenceSection
-
-
-class SearchAnswerResults(BaseModel):
-    query: str
-    answer: str
-    quality: str
-    documents: Annotated[list[InferenceSection], dedup_inference_sections]
-
-
-class QACheckOutput(TypedDict, total=False):
-    answer_quality: str
-
-
-class QAGenerationOutput(TypedDict, total=False):
-    answer: str
-
-
-class ExpandedRetrievalOutput(TypedDict):
-    reordered_documents: Annotated[list[InferenceSection], dedup_inference_sections]
-
-
-class AnswerQueryState(
-    PrimaryState,
-    QACheckOutput,
-    QAGenerationOutput,
-    ExpandedRetrievalOutput,
-    total=True,
-):
-    query_to_answer: str
-
-
-class AnswerQueryInput(PrimaryState, total=True):
-    query_to_answer: str
-
-
-class AnswerQueryOutput(TypedDict):
-    decomp_answer_results: list[SearchAnswerResults]

backend/danswer/agent_search/core_state.py

@@ -1,15 +0,0 @@
-from typing import TypedDict
-
-from sqlalchemy.orm import Session
-
-from danswer.context.search.models import SearchRequest
-from danswer.llm.interfaces import LLM
-
-
-class PrimaryState(TypedDict, total=False):
-    search_request: SearchRequest
-    primary_llm: LLM
-    fast_llm: LLM
-    # a single session for the entire agent search
-    # is fine if we are only reading
-    db_session: Session

backend/danswer/agent_search/deep_answer/nodes/answer_generation.py

@@ -1,114 +0,0 @@
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.main.states import MainState
-from danswer.agent_search.shared_graph_utils.prompts import COMBINED_CONTEXT
-from danswer.agent_search.shared_graph_utils.prompts import MODIFIED_RAG_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-from danswer.agent_search.shared_graph_utils.utils import normalize_whitespace
-
-
-# aggregate sub questions and answers
-def deep_answer_generation(state: MainState) -> dict[str, Any]:
-    """
-    Generate answer
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-         dict: The updated state with re-phrased question
-    """
-    print("---DEEP GENERATE---")
-
-    question = state["original_question"]
-    docs = state["deduped_retrieval_docs"]
-
-    deep_answer_context = state["core_answer_dynamic_context"]
-
-    print(f"Number of verified retrieval docs - deep: {len(docs)}")
-
-    combined_context = normalize_whitespace(
-        COMBINED_CONTEXT.format(
-            deep_answer_context=deep_answer_context, formated_docs=format_docs(docs)
-        )
-    )
-
-    msg = [
-        HumanMessage(
-            content=MODIFIED_RAG_PROMPT.format(
-                question=question, combined_context=combined_context
-            )
-        )
-    ]
-
-    # Grader
-    model = state["fast_llm"]
-    response = model.invoke(msg)
-
-    return {
-        "deep_answer": response.content,
-    }
-
-
-def final_stuff(state: MainState) -> dict[str, Any]:
-    """
-    Invokes the agent model to generate a response based on the current state. Given
-    the question, it will decide to retrieve using the retriever tool, or simply end.
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-        dict: The updated state with the agent response appended to messages
-    """
-    print("---FINAL---")
-
-    messages = state["log_messages"]
-    time_ordered_messages = [x.pretty_repr() for x in messages]
-    time_ordered_messages.sort()
-
-    print("Message Log:")
-    print("\n".join(time_ordered_messages))
-
-    initial_sub_qas = state["initial_sub_qas"]
-    initial_sub_qa_list = []
-    for initial_sub_qa in initial_sub_qas:
-        if initial_sub_qa["sub_answer_check"] == "yes":
-            initial_sub_qa_list.append(
-                f'  Question:\n  {initial_sub_qa["sub_question"]}\n  --\n  Answer:\n  {initial_sub_qa["sub_answer"]}\n  -----'
-            )
-
-    initial_sub_qa_context = "\n".join(initial_sub_qa_list)
-
-    base_answer = state["base_answer"]
-
-    print(f"Final Base Answer:\n{base_answer}")
-    print("--------------------------------")
-    print(f"Initial Answered Sub Questions:\n{initial_sub_qa_context}")
-    print("--------------------------------")
-
-    if not state.get("deep_answer"):
-        print("No Deep Answer was required")
-        return {}
-
-    deep_answer = state["deep_answer"]
-    sub_qas = state["sub_qas"]
-    sub_qa_list = []
-    for sub_qa in sub_qas:
-        if sub_qa["sub_answer_check"] == "yes":
-            sub_qa_list.append(
-                f'  Question:\n  {sub_qa["sub_question"]}\n  --\n  Answer:\n  {sub_qa["sub_answer"]}\n  -----'
-            )
-
-    sub_qa_context = "\n".join(sub_qa_list)
-
-    print(f"Final Base Answer:\n{base_answer}")
-    print("--------------------------------")
-    print(f"Final Deep Answer:\n{deep_answer}")
-    print("--------------------------------")
-    print("Sub Questions and Answers:")
-    print(sub_qa_context)
-
-    return {}

backend/danswer/agent_search/deep_answer/nodes/deep_decomp.py

@@ -1,78 +0,0 @@
-import json
-import re
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.main.states import MainState
-from danswer.agent_search.shared_graph_utils.prompts import DEEP_DECOMPOSE_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_entity_term_extraction
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def decompose(state: MainState) -> dict[str, Any]:
-    """ """
-
-    node_start_time = datetime.now()
-
-    question = state["original_question"]
-    base_answer = state["base_answer"]
-
-    # get the entity term extraction dict and properly format it
-    entity_term_extraction_dict = state["retrieved_entities_relationships"][
-        "retrieved_entities_relationships"
-    ]
-
-    entity_term_extraction_str = format_entity_term_extraction(
-        entity_term_extraction_dict
-    )
-
-    initial_question_answers = state["initial_sub_qas"]
-
-    addressed_question_list = [
-        x["sub_question"]
-        for x in initial_question_answers
-        if x["sub_answer_check"] == "yes"
-    ]
-    failed_question_list = [
-        x["sub_question"]
-        for x in initial_question_answers
-        if x["sub_answer_check"] == "no"
-    ]
-
-    msg = [
-        HumanMessage(
-            content=DEEP_DECOMPOSE_PROMPT.format(
-                question=question,
-                entity_term_extraction_str=entity_term_extraction_str,
-                base_answer=base_answer,
-                answered_sub_questions="\n - ".join(addressed_question_list),
-                failed_sub_questions="\n - ".join(failed_question_list),
-            ),
-        )
-    ]
-
-    # Grader
-    model = state["fast_llm"]
-    response = model.invoke(msg)
-
-    cleaned_response = re.sub(r"```json\n|\n```", "", response.pretty_repr())
-    parsed_response = json.loads(cleaned_response)
-
-    sub_questions_dict = {}
-    for sub_question_nr, sub_question_dict in enumerate(
-        parsed_response["sub_questions"]
-    ):
-        sub_question_dict["answered"] = False
-        sub_question_dict["verified"] = False
-        sub_questions_dict[sub_question_nr] = sub_question_dict
-
-    return {
-        "decomposed_sub_questions_dict": sub_questions_dict,
-        "log_messages": generate_log_message(
-            message="deep - decompose",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_answer/nodes/entity_term_extraction.py

@@ -1,40 +0,0 @@
-import json
-import re
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.main.states import MainState
-from danswer.agent_search.shared_graph_utils.prompts import ENTITY_TERM_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-
-
-def entity_term_extraction(state: MainState) -> dict[str, Any]:
-    """Extract entities and terms from the question and context"""
-
-    question = state["original_question"]
-    docs = state["deduped_retrieval_docs"]
-
-    doc_context = format_docs(docs)
-
-    msg = [
-        HumanMessage(
-            content=ENTITY_TERM_PROMPT.format(question=question, context=doc_context),
-        )
-    ]
-    fast_llm = state["fast_llm"]
-    # Grader
-    llm_response_list = list(
-        fast_llm.stream(
-            prompt=msg,
-        )
-    )
-    llm_response = merge_message_runs(llm_response_list, chunk_separator="")[0].content
-
-    cleaned_response = re.sub(r"```json\n|\n```", "", llm_response)
-    parsed_response = json.loads(cleaned_response)
-
-    return {
-        "retrieved_entities_relationships": parsed_response,
-    }

backend/danswer/agent_search/deep_answer/nodes/sub_qa_level_aggregator.py

@@ -1,30 +0,0 @@
-from typing import Any
-
-from danswer.agent_search.main.states import MainState
-
-
-# aggregate sub questions and answers
-def sub_qa_level_aggregator(state: MainState) -> dict[str, Any]:
-    sub_qas = state["sub_qas"]
-
-    dynamic_context_list = [
-        "Below you will find useful information to answer the original question:"
-    ]
-    checked_sub_qas = []
-
-    for core_answer_sub_qa in sub_qas:
-        question = core_answer_sub_qa["sub_question"]
-        answer = core_answer_sub_qa["sub_answer"]
-        verified = core_answer_sub_qa["sub_answer_check"]
-
-        if verified == "yes":
-            dynamic_context_list.append(
-                f"Question:\n{question}\n\nAnswer:\n{answer}\n\n---\n\n"
-            )
-            checked_sub_qas.append({"sub_question": question, "sub_answer": answer})
-    dynamic_context = "\n".join(dynamic_context_list)
-
-    return {
-        "core_answer_dynamic_context": dynamic_context,
-        "checked_sub_qas": checked_sub_qas,
-    }

backend/danswer/agent_search/deep_answer/nodes/sub_qa_manager.py

@@ -1,19 +0,0 @@
-from typing import Any
-
-from danswer.agent_search.main.states import MainState
-
-
-def sub_qa_manager(state: MainState) -> dict[str, Any]:
-    """ """
-
-    sub_questions_dict = state["decomposed_sub_questions_dict"]
-
-    sub_questions = {}
-
-    for sub_question_nr, sub_question_dict in sub_questions_dict.items():
-        sub_questions[sub_question_nr] = sub_question_dict["sub_question"]
-
-    return {
-        "sub_questions": sub_questions,
-        "num_new_question_iterations": 0,
-    }

backend/danswer/agent_search/expanded_retrieval/edges.py

@@ -1,44 +0,0 @@
-from collections.abc import Hashable
-
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-from langgraph.types import Send
-
-from danswer.agent_search.expanded_retrieval.nodes.doc_retrieval import RetrieveInput
-from danswer.agent_search.expanded_retrieval.states import ExpandedRetrievalInput
-from danswer.agent_search.shared_graph_utils.prompts import REWRITE_PROMPT_MULTI
-from danswer.llm.interfaces import LLM
-
-
-def parallel_retrieval_edge(state: ExpandedRetrievalInput) -> list[Send | Hashable]:
-    print(f"parallel_retrieval_edge state: {state.keys()}")
-
-    # This should be better...
-    question = state.get("query_to_answer") or state["search_request"].query
-    llm: LLM = state["fast_llm"]
-
-    msg = [
-        HumanMessage(
-            content=REWRITE_PROMPT_MULTI.format(question=question),
-        )
-    ]
-    llm_response_list = list(
-        llm.stream(
-            prompt=msg,
-        )
-    )
-    llm_response = merge_message_runs(llm_response_list, chunk_separator="")[0].content
-
-    print(f"llm_response: {llm_response}")
-
-    rewritten_queries = llm_response.split("\n")
-
-    print(f"rewritten_queries: {rewritten_queries}")
-
-    return [
-        Send(
-            "doc_retrieval",
-            RetrieveInput(query_to_retrieve=query, **state),
-        )
-        for query in rewritten_queries
-    ]

backend/danswer/agent_search/expanded_retrieval/graph_builder.py

@@ -1,88 +0,0 @@
-from langgraph.graph import END
-from langgraph.graph import START
-from langgraph.graph import StateGraph
-
-from danswer.agent_search.expanded_retrieval.edges import parallel_retrieval_edge
-from danswer.agent_search.expanded_retrieval.nodes.doc_reranking import doc_reranking
-from danswer.agent_search.expanded_retrieval.nodes.doc_retrieval import doc_retrieval
-from danswer.agent_search.expanded_retrieval.nodes.doc_verification import (
-    doc_verification,
-)
-from danswer.agent_search.expanded_retrieval.nodes.verification_kickoff import (
-    verification_kickoff,
-)
-from danswer.agent_search.expanded_retrieval.states import ExpandedRetrievalInput
-from danswer.agent_search.expanded_retrieval.states import ExpandedRetrievalOutput
-from danswer.agent_search.expanded_retrieval.states import ExpandedRetrievalState
-
-
-def expanded_retrieval_graph_builder() -> StateGraph:
-    graph = StateGraph(
-        state_schema=ExpandedRetrievalState,
-        input=ExpandedRetrievalInput,
-        output=ExpandedRetrievalOutput,
-    )
-
-    ### Add nodes ###
-
-    graph.add_node(
-        node="doc_retrieval",
-        action=doc_retrieval,
-    )
-    graph.add_node(
-        node="verification_kickoff",
-        action=verification_kickoff,
-    )
-    graph.add_node(
-        node="doc_verification",
-        action=doc_verification,
-    )
-    graph.add_node(
-        node="doc_reranking",
-        action=doc_reranking,
-    )
-
-    ### Add edges ###
-
-    graph.add_conditional_edges(
-        source=START,
-        path=parallel_retrieval_edge,
-        path_map=["doc_retrieval"],
-    )
-    graph.add_edge(
-        start_key="doc_retrieval",
-        end_key="verification_kickoff",
-    )
-    graph.add_edge(
-        start_key="doc_verification",
-        end_key="doc_reranking",
-    )
-    graph.add_edge(
-        start_key="doc_reranking",
-        end_key=END,
-    )
-
-    return graph
-
-
-if __name__ == "__main__":
-    from danswer.db.engine import get_session_context_manager
-    from danswer.llm.factory import get_default_llms
-    from danswer.context.search.models import SearchRequest
-
-    graph = expanded_retrieval_graph_builder()
-    compiled_graph = graph.compile()
-    primary_llm, fast_llm = get_default_llms()
-    search_request = SearchRequest(
-        query="Who made Excel and what other products did they make?",
-    )
-    with get_session_context_manager() as db_session:
-        inputs = ExpandedRetrievalInput(
-            search_request=search_request,
-            primary_llm=primary_llm,
-            fast_llm=fast_llm,
-            db_session=db_session,
-            query_to_answer="Who made Excel?",
-        )
-        for thing in compiled_graph.stream(inputs, debug=True):
-            print(thing)

backend/danswer/agent_search/expanded_retrieval/nodes/doc_reranking.py

@@ -1,11 +0,0 @@
-from danswer.agent_search.expanded_retrieval.states import DocRerankingOutput
-from danswer.agent_search.expanded_retrieval.states import ExpandedRetrievalState
-
-
-def doc_reranking(state: ExpandedRetrievalState) -> DocRerankingOutput:
-    print(f"doc_reranking state: {state.keys()}")
-
-    verified_documents = state["verified_documents"]
-    reranked_documents = verified_documents
-
-    return DocRerankingOutput(reranked_documents=reranked_documents)

backend/danswer/agent_search/expanded_retrieval/nodes/doc_retrieval.py

@@ -1,47 +0,0 @@
-from danswer.agent_search.expanded_retrieval.states import DocRetrievalOutput
-from danswer.agent_search.expanded_retrieval.states import ExpandedRetrievalState
-from danswer.context.search.models import InferenceSection
-from danswer.context.search.models import SearchRequest
-from danswer.context.search.pipeline import SearchPipeline
-from danswer.db.engine import get_session_context_manager
-
-
-class RetrieveInput(ExpandedRetrievalState):
-    query_to_retrieve: str
-
-
-def doc_retrieval(state: RetrieveInput) -> DocRetrievalOutput:
-    # def doc_retrieval(state: RetrieveInput) -> Command[Literal["doc_verification"]]:
-    """
-    Retrieve documents
-
-    Args:
-        state (dict): The current graph state
-
-    Returns:
-        state (dict): New key added to state, documents, that contains retrieved documents
-    """
-    print(f"doc_retrieval state: {state.keys()}")
-
-    state["query_to_retrieve"]
-
-    documents: list[InferenceSection] = []
-    llm = state["primary_llm"]
-    fast_llm = state["fast_llm"]
-    # db_session = state["db_session"]
-    query_to_retrieve = state["search_request"].query
-    with get_session_context_manager() as db_session1:
-        documents = SearchPipeline(
-            search_request=SearchRequest(
-                query=query_to_retrieve,
-            ),
-            user=None,
-            llm=llm,
-            fast_llm=fast_llm,
-            db_session=db_session1,
-        ).reranked_sections
-
-    print(f"retrieved documents: {len(documents)}")
-    return DocRetrievalOutput(
-        retrieved_documents=documents,
-    )

backend/danswer/agent_search/expanded_retrieval/nodes/doc_verification.py

@@ -1,60 +0,0 @@
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.expanded_retrieval.states import DocVerificationOutput
-from danswer.agent_search.expanded_retrieval.states import ExpandedRetrievalState
-from danswer.agent_search.shared_graph_utils.models import BinaryDecision
-from danswer.agent_search.shared_graph_utils.prompts import VERIFIER_PROMPT
-from danswer.context.search.models import InferenceSection
-
-
-class DocVerificationInput(ExpandedRetrievalState, total=True):
-    doc_to_verify: InferenceSection
-
-
-def doc_verification(state: DocVerificationInput) -> DocVerificationOutput:
-    """
-    Check whether the document is relevant for the original user question
-
-    Args:
-        state (VerifierState): The current state
-
-    Returns:
-        dict: ict: The updated state with the final decision
-    """
-
-    print(f"doc_verification state: {state.keys()}")
-
-    original_query = state["search_request"].query
-    doc_to_verify = state["doc_to_verify"]
-    document_content = doc_to_verify.combined_content
-
-    msg = [
-        HumanMessage(
-            content=VERIFIER_PROMPT.format(
-                question=original_query, document_content=document_content
-            )
-        )
-    ]
-
-    fast_llm = state["fast_llm"]
-    response = list(
-        fast_llm.stream(
-            prompt=msg,
-        )
-    )
-
-    response_string = merge_message_runs(response, chunk_separator="")[0].content
-    # Convert string response to proper dictionary format
-    decision_dict = {"decision": response_string.lower()}
-    formatted_response = BinaryDecision.model_validate(decision_dict)
-
-    print(f"Verdict: {formatted_response.decision}")
-
-    verified_documents = []
-    if formatted_response.decision == "yes":
-        verified_documents.append(doc_to_verify)
-
-    return DocVerificationOutput(
-        verified_documents=verified_documents,
-    )

backend/danswer/agent_search/expanded_retrieval/nodes/verification_kickoff.py

@@ -1,27 +0,0 @@
-from typing import Literal
-
-from langgraph.types import Command
-from langgraph.types import Send
-
-from danswer.agent_search.expanded_retrieval.nodes.doc_verification import (
-    DocVerificationInput,
-)
-from danswer.agent_search.expanded_retrieval.states import ExpandedRetrievalState
-
-
-def verification_kickoff(
-    state: ExpandedRetrievalState,
-) -> Command[Literal["doc_verification"]]:
-    print(f"verification_kickoff state: {state.keys()}")
-
-    documents = state["retrieved_documents"]
-    return Command(
-        update={},
-        goto=[
-            Send(
-                node="doc_verification",
-                arg=DocVerificationInput(doc_to_verify=doc, **state),
-            )
-            for doc in documents
-        ],
-    )

backend/danswer/agent_search/expanded_retrieval/states.py

@@ -1,36 +0,0 @@
-from typing import Annotated
-from typing import TypedDict
-
-from danswer.agent_search.core_state import PrimaryState
-from danswer.agent_search.shared_graph_utils.operators import dedup_inference_sections
-from danswer.context.search.models import InferenceSection
-
-
-class DocRetrievalOutput(TypedDict, total=False):
-    retrieved_documents: Annotated[list[InferenceSection], dedup_inference_sections]
-
-
-class DocVerificationOutput(TypedDict, total=False):
-    verified_documents: Annotated[list[InferenceSection], dedup_inference_sections]
-
-
-class DocRerankingOutput(TypedDict, total=False):
-    reranked_documents: Annotated[list[InferenceSection], dedup_inference_sections]
-
-
-class ExpandedRetrievalState(
-    PrimaryState,
-    DocRetrievalOutput,
-    DocVerificationOutput,
-    DocRerankingOutput,
-    total=True,
-):
-    query_to_answer: str
-
-
-class ExpandedRetrievalInput(PrimaryState, total=True):
-    query_to_answer: str
-
-
-class ExpandedRetrievalOutput(TypedDict):
-    reordered_documents: Annotated[list[InferenceSection], dedup_inference_sections]

backend/danswer/agent_search/main/edges.py

@@ -1,61 +0,0 @@
-from collections.abc import Hashable
-
-from langgraph.types import Send
-
-from danswer.agent_search.answer_query.states import AnswerQueryInput
-from danswer.agent_search.main.states import MainState
-
-
-def parallelize_decompozed_answer_queries(state: MainState) -> list[Send | Hashable]:
-    return [
-        Send(
-            "answer_query",
-            AnswerQueryInput(
-                **state,
-                query_to_answer=query,
-            ),
-        )
-        for query in state["initial_decomp_queries"]
-    ]
-
-
-# def continue_to_answer_sub_questions(state: QAState) -> Union[Hashable, list[Hashable]]:
-#     # Routes re-written queries to the (parallel) retrieval steps
-#     # Notice the 'Send()' API that takes care of the parallelization
-#     return [
-#         Send(
-#             "sub_answers_graph",
-#             ResearchQAState(
-#                 sub_question=sub_question["sub_question_str"],
-#                 sub_question_nr=sub_question["sub_question_nr"],
-#                 graph_start_time=state["graph_start_time"],
-#                 primary_llm=state["primary_llm"],
-#                 fast_llm=state["fast_llm"],
-#             ),
-#         )
-#         for sub_question in state["sub_questions"]
-#     ]
-
-
-# def continue_to_deep_answer(state: QAState) -> Union[Hashable, list[Hashable]]:
-#     print("---GO TO DEEP ANSWER OR END---")
-
-#     base_answer = state["base_answer"]
-
-#     question = state["original_question"]
-
-#     BASE_CHECK_MESSAGE = [
-#         HumanMessage(
-#             content=BASE_CHECK_PROMPT.format(question=question, base_answer=base_answer)
-#         )
-#     ]
-
-#     model = state["fast_llm"]
-#     response = model.invoke(BASE_CHECK_MESSAGE)
-
-#     print(f"CAN WE CONTINUE W/O GENERATING A DEEP ANSWER? - {response.pretty_repr()}")
-
-#     if response.pretty_repr() == "no":
-#         return "decompose"
-#     else:
-#         return "end"

backend/danswer/agent_search/main/graph_builder.py

@@ -1,98 +0,0 @@
-from langgraph.graph import END
-from langgraph.graph import START
-from langgraph.graph import StateGraph
-
-from danswer.agent_search.answer_query.graph_builder import answer_query_graph_builder
-from danswer.agent_search.expanded_retrieval.graph_builder import (
-    expanded_retrieval_graph_builder,
-)
-from danswer.agent_search.main.edges import parallelize_decompozed_answer_queries
-from danswer.agent_search.main.nodes.base_decomp import main_decomp_base
-from danswer.agent_search.main.nodes.generate_initial_answer import (
-    generate_initial_answer,
-)
-from danswer.agent_search.main.states import MainInput
-from danswer.agent_search.main.states import MainState
-
-
-def main_graph_builder() -> StateGraph:
-    graph = StateGraph(
-        state_schema=MainState,
-        input=MainInput,
-    )
-
-    ### Add nodes ###
-
-    graph.add_node(
-        node="base_decomp",
-        action=main_decomp_base,
-    )
-    answer_query_subgraph = answer_query_graph_builder().compile()
-    graph.add_node(
-        node="answer_query",
-        action=answer_query_subgraph,
-    )
-    expanded_retrieval_subgraph = expanded_retrieval_graph_builder().compile()
-    graph.add_node(
-        node="expanded_retrieval",
-        action=expanded_retrieval_subgraph,
-    )
-    graph.add_node(
-        node="generate_initial_answer",
-        action=generate_initial_answer,
-    )
-
-    ### Add edges ###
-    graph.add_edge(
-        start_key=START,
-        end_key="expanded_retrieval",
-    )
-
-    graph.add_edge(
-        start_key=START,
-        end_key="base_decomp",
-    )
-    graph.add_conditional_edges(
-        source="base_decomp",
-        path=parallelize_decompozed_answer_queries,
-        path_map=["answer_query"],
-    )
-    graph.add_edge(
-        start_key=["answer_query", "expanded_retrieval"],
-        end_key="generate_initial_answer",
-    )
-    graph.add_edge(
-        start_key="generate_initial_answer",
-        end_key=END,
-    )
-
-    return graph
-
-
-if __name__ == "__main__":
-    from danswer.db.engine import get_session_context_manager
-    from danswer.llm.factory import get_default_llms
-    from danswer.context.search.models import SearchRequest
-
-    graph = main_graph_builder()
-    compiled_graph = graph.compile()
-    primary_llm, fast_llm = get_default_llms()
-    search_request = SearchRequest(
-        query="If i am familiar with the function that I need, how can I type it into a cell?",
-    )
-    with get_session_context_manager() as db_session:
-        inputs = MainInput(
-            search_request=search_request,
-            primary_llm=primary_llm,
-            fast_llm=fast_llm,
-            db_session=db_session,
-        )
-        for thing in compiled_graph.stream(
-            input=inputs,
-            # stream_mode="debug",
-            # debug=True,
-            subgraphs=True,
-        ):
-            # print(thing)
-            print()
-            print()

backend/danswer/agent_search/main/nodes/base_decomp.py

@@ -1,31 +0,0 @@
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.main.states import BaseDecompOutput
-from danswer.agent_search.main.states import MainState
-from danswer.agent_search.shared_graph_utils.prompts import INITIAL_DECOMPOSITION_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import clean_and_parse_list_string
-
-
-def main_decomp_base(state: MainState) -> BaseDecompOutput:
-    question = state["search_request"].query
-
-    msg = [
-        HumanMessage(
-            content=INITIAL_DECOMPOSITION_PROMPT.format(question=question),
-        )
-    ]
-
-    # Get the rewritten queries in a defined format
-    model = state["fast_llm"]
-    response = model.invoke(msg)
-
-    content = response.pretty_repr()
-    list_of_subquestions = clean_and_parse_list_string(content)
-
-    decomp_list: list[str] = [
-        sub_question["sub_question"].strip() for sub_question in list_of_subquestions
-    ]
-
-    return BaseDecompOutput(
-        initial_decomp_queries=decomp_list,
-    )

backend/danswer/agent_search/main/nodes/generate_initial_answer.py

@@ -1,53 +0,0 @@
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.main.states import InitialAnswerOutput
-from danswer.agent_search.main.states import MainState
-from danswer.agent_search.shared_graph_utils.prompts import INITIAL_RAG_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-
-
-def generate_initial_answer(state: MainState) -> InitialAnswerOutput:
-    print("---GENERATE INITIAL---")
-
-    question = state["search_request"].query
-    docs = state["documents"]
-
-    decomp_answer_results = state["decomp_answer_results"]
-
-    good_qa_list: list[str] = []
-
-    _SUB_QUESTION_ANSWER_TEMPLATE = """
-    Sub-Question:\n  - {sub_question}\n  --\nAnswer:\n  - {sub_answer}\n\n
-    """
-    for decomp_answer_result in decomp_answer_results:
-        if (
-            decomp_answer_result.quality.lower() == "yes"
-            and len(decomp_answer_result.answer) > 0
-            and decomp_answer_result.answer != "I don't know"
-        ):
-            good_qa_list.append(
-                _SUB_QUESTION_ANSWER_TEMPLATE.format(
-                    sub_question=decomp_answer_result.query,
-                    sub_answer=decomp_answer_result.answer,
-                )
-            )
-
-    sub_question_answer_str = "\n\n------\n\n".join(good_qa_list)
-
-    msg = [
-        HumanMessage(
-            content=INITIAL_RAG_PROMPT.format(
-                question=question,
-                context=format_docs(docs),
-                answered_sub_questions=sub_question_answer_str,
-            )
-        )
-    ]
-
-    # Grader
-    model = state["fast_llm"]
-    response = model.invoke(msg)
-    answer = response.pretty_repr()
-
-    print(answer)
-    return InitialAnswerOutput(initial_answer=answer)

backend/danswer/agent_search/main/states.py

@@ -1,37 +0,0 @@
-from operator import add
-from typing import Annotated
-from typing import TypedDict
-
-from danswer.agent_search.answer_query.states import SearchAnswerResults
-from danswer.agent_search.core_state import PrimaryState
-from danswer.agent_search.shared_graph_utils.operators import dedup_inference_sections
-from danswer.context.search.models import InferenceSection
-
-
-class BaseDecompOutput(TypedDict, total=False):
-    initial_decomp_queries: list[str]
-
-
-class InitialAnswerOutput(TypedDict, total=False):
-    initial_answer: str
-
-
-class MainState(
-    PrimaryState,
-    BaseDecompOutput,
-    InitialAnswerOutput,
-    total=True,
-):
-    documents: Annotated[list[InferenceSection], dedup_inference_sections]
-    decomp_answer_results: Annotated[list[SearchAnswerResults], add]
-
-
-class MainInput(PrimaryState, total=True):
-    pass
-
-
-class MainOutput(TypedDict):
-    """
-    This is not used because defining the output only matters for filtering the output of
-      a .invoke() call but we are streaming so we just yield the entire state.
-    """

backend/danswer/agent_search/run_graph.py

@@ -1,27 +0,0 @@
-from danswer.agent_search.primary_graph.graph_builder import build_core_graph
-from danswer.llm.answering.answer import AnswerStream
-from danswer.llm.interfaces import LLM
-from danswer.tools.tool import Tool
-
-
-def run_graph(
-    query: str,
-    llm: LLM,
-    tools: list[Tool],
-) -> AnswerStream:
-    graph = build_core_graph()
-
-    inputs = {
-        "original_query": query,
-        "messages": [],
-        "tools": tools,
-        "llm": llm,
-    }
-    compiled_graph = graph.compile()
-    output = compiled_graph.invoke(input=inputs)
-    yield from output
-
-
-if __name__ == "__main__":
-    pass
-    # run_graph("What is the capital of France?", llm, [])

backend/danswer/agent_search/shared_graph_utils/models.py

@@ -1,12 +0,0 @@
-from typing import Literal
-
-from pydantic import BaseModel
-
-
-# Pydantic models for structured outputs
-class RewrittenQueries(BaseModel):
-    rewritten_queries: list[str]
-
-
-class BinaryDecision(BaseModel):
-    decision: Literal["yes", "no"]

backend/danswer/agent_search/shared_graph_utils/operators.py

@@ -1,9 +0,0 @@
-from danswer.context.search.models import InferenceSection
-from danswer.llm.answering.prune_and_merge import _merge_sections
-
-
-def dedup_inference_sections(
-    list1: list[InferenceSection], list2: list[InferenceSection]
-) -> list[InferenceSection]:
-    deduped = _merge_sections(list1 + list2)
-    return deduped

backend/danswer/agent_search/shared_graph_utils/prompts.py

@@ -1,427 +0,0 @@
-REWRITE_PROMPT_MULTI_ORIGINAL = """ \n
-    Please convert an initial user question into a 2-3 more appropriate short and pointed search queries for retrievel from a
-    document store. Particularly, try to think about resolving ambiguities and make the search queries more specific,
-    enabling the system to search more broadly.
-    Also, try to make the search queries not redundant, i.e. not too similar! \n\n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-    Formulate the queries separated by '--' (Do not say 'Query 1: ...', just write the querytext): """
-
-REWRITE_PROMPT_MULTI = """ \n
-    Please create a list of 2-3 sample documents that could answer an original question. Each document
-    should be about as long as the original question. \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-    Formulate the sample documents separated by '--' (Do not say 'Document 1: ...', just write the text): """
-
-BASE_RAG_PROMPT = """ \n
-    You are an assistant for question-answering tasks. Use the context provided below - and only the
-    provided context - to answer the question. If you don't know the answer or if the provided context is
-    empty, just say "I don't know". Do not use your internal knowledge!
-
-    Again, only use the provided context and do not use your internal knowledge! If you cannot answer the
-    question based on the context, say "I don't know". It is a matter of life and death that you do NOT
-    use your internal knowledge, just the provided information!
-
-    Use three sentences maximum and keep the answer concise.
-    answer concise.\nQuestion:\n {question} \nContext:\n {context} \n\n
-    \n\n
-    Answer:"""
-
-BASE_CHECK_PROMPT = """ \n
-    Please check whether 1) the suggested answer seems to fully address the original question AND 2)the
-    original question requests a simple, factual answer, and there are no ambiguities, judgements,
-    aggregations, or any other complications that may require extra context. (I.e., if the question is
-    somewhat addressed, but the answer would benefit from more context, then answer with 'no'.)
-
-    Please only answer with 'yes' or 'no' \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-    Here is the proposed answer:
-    \n ------- \n
-    {base_answer}
-    \n ------- \n
-    Please answer with yes or no:"""
-
-VERIFIER_PROMPT = """ \n
-    Please check whether the document seems to be relevant for the answer of the question. Please
-    only answer with 'yes' or 'no' \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-    Here is the document text:
-    \n ------- \n
-    {document_content}
-    \n ------- \n
-    Please answer with yes or no:"""
-
-INITIAL_DECOMPOSITION_PROMPT_BASIC = """ \n
-    Please decompose an initial user question into not more than 4 appropriate sub-questions that help to
-    answer the original question. The purpose for this decomposition is to isolate individulal entities
-    (i.e., 'compare sales of company A and company B' -> 'what are sales for company A' + 'what are sales
-    for company B'), split ambiguous terms (i.e., 'what is our success with company A' -> 'what are our
-    sales with company A' + 'what is our market share with company A' + 'is company A a reference customer
-    for us'), etc. Each sub-question should be realistically be answerable by a good RAG system. \n
-
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Please formulate your answer as a list of subquestions:
-
-    Answer:
-    """
-
-REWRITE_PROMPT_SINGLE = """ \n
-    Please convert an initial user question into a more appropriate search query for retrievel from a
-    document store. \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Formulate the query: """
-
-MODIFIED_RAG_PROMPT = """You are an assistant for question-answering tasks. Use the context provided below
-    - and only this context - to answer the question. If you don't know the answer, just say "I don't know".
-    Use three sentences maximum and keep the answer concise.
-    Pay also particular attention to the sub-questions and their answers, at least it may enrich the answer.
-    Again, only use the provided context and do not use your internal knowledge! If you cannot answer the
-    question based on the context, say "I don't know". It is a matter of life and death that you do NOT
-    use your internal knowledge, just the provided information!
-
-    \nQuestion: {question}
-    \nContext: {combined_context} \n
-
-    Answer:"""
-
-ORIG_DEEP_DECOMPOSE_PROMPT = """ \n
-    An initial user question needs to be answered. An initial answer has been provided but it wasn't quite
-    good enough. Also, some sub-questions had been answered and this information has been used to provide
-    the initial answer. Some other subquestions may have been suggested based on little knowledge, but they
-    were not directly answerable. Also, some entities, relationships and terms are givenm to you so that
-    you have an idea of how the avaiolable data looks like.
-
-    Your role is to generate 3-5 new sub-questions that would help to answer the initial question,
-    considering:
-
-    1) The initial question
-    2) The initial answer that was found to be unsatisfactory
-    3) The sub-questions that were answered
-    4) The sub-questions that were suggested but not answered
-    5) The entities, relationships and terms that were extracted from the context
-
-    The individual questions should be answerable by a good RAG system.
-    So a good idea would be to use the sub-questions to resolve ambiguities and/or to separate the
-    question for different entities that may be involved in the original question, but in a way that does
-    not duplicate questions that were already tried.
-
-    Additional Guidelines:
-    - The sub-questions should be specific to the question and provide richer context for the question,
-    resolve ambiguities, or address shortcoming of the initial answer
-    - Each sub-question - when answered - should be relevant for the answer to the original question
-    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
-    other complications that may require extra context.
-    - The sub-questions MUST have the full context of the original question so that it can be executed by
-    a RAG system independently without the original question available
-      (Example:
-        - initial question: "What is the capital of France?"
-        - bad sub-question: "What is the name of the river there?"
-        - good sub-question: "What is the name of the river that flows through Paris?"
-    - For each sub-question, please provide a short explanation for why it is a good sub-question. So
-    generate a list of dictionaries with the following format:
-      [{{"sub_question": <sub-question>, "explanation": <explanation>, "search_term": <rewrite the
-      sub-question using as a search phrase for the document store>}}, ...]
-
-    \n\n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Here is the initial sub-optimal answer:
-    \n ------- \n
-    {base_answer}
-    \n ------- \n
-
-    Here are the sub-questions that were answered:
-    \n ------- \n
-    {answered_sub_questions}
-    \n ------- \n
-
-    Here are the sub-questions that were suggested but not answered:
-    \n ------- \n
-    {failed_sub_questions}
-    \n ------- \n
-
-    And here are the entities, relationships and terms extracted from the context:
-    \n ------- \n
-    {entity_term_extraction_str}
-    \n ------- \n
-
-   Please generate the list of good, fully contextualized sub-questions that would help to address the
-   main question. Again, please find questions that are NOT overlapping too much with the already answered
-   sub-questions or those that already were suggested and failed.
-   In other words - what can we try in addition to what has been tried so far?
-
-   Please think through it step by step and then generate the list of json dictionaries with the following
-   format:
-
-   {{"sub_questions": [{{"sub_question": <sub-question>,
-        "explanation": <explanation>,
-        "search_term": <rewrite the sub-question using as a search phrase for the document store>}},
-        ...]}} """
-
-DEEP_DECOMPOSE_PROMPT = """ \n
-    An initial user question needs to be answered. An initial answer has been provided but it wasn't quite
-    good enough. Also, some sub-questions had been answered and this information has been used to provide
-    the initial answer. Some other subquestions may have been suggested based on little knowledge, but they
-    were not directly answerable. Also, some entities, relationships and terms are givenm to you so that
-    you have an idea of how the avaiolable data looks like.
-
-    Your role is to generate 4-6 new sub-questions that would help to answer the initial question,
-    considering:
-
-    1) The initial question
-    2) The initial answer that was found to be unsatisfactory
-    3) The sub-questions that were answered
-    4) The sub-questions that were suggested but not answered
-    5) The entities, relationships and terms that were extracted from the context
-
-    The individual questions should be answerable by a good RAG system.
-    So a good idea would be to use the sub-questions to resolve ambiguities and/or to separate the
-    question for different entities that may be involved in the original question, but in a way that does
-    not duplicate questions that were already tried.
-
-    Additional Guidelines:
-    - The sub-questions should be specific to the question and provide richer context for the question,
-    resolve ambiguities, or address shortcoming of the initial answer
-    - Each sub-question - when answered - should be relevant for the answer to the original question
-    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
-    other complications that may require extra context.
-    - The sub-questions MUST have the full context of the original question so that it can be executed by
-    a RAG system independently without the original question available
-      (Example:
-        - initial question: "What is the capital of France?"
-        - bad sub-question: "What is the name of the river there?"
-        - good sub-question: "What is the name of the river that flows through Paris?"
-    - For each sub-question, please also provide a search term that can be used to retrieve relevant
-    documents from a document store.
-    \n\n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Here is the initial sub-optimal answer:
-    \n ------- \n
-    {base_answer}
-    \n ------- \n
-
-    Here are the sub-questions that were answered:
-    \n ------- \n
-    {answered_sub_questions}
-    \n ------- \n
-
-    Here are the sub-questions that were suggested but not answered:
-    \n ------- \n
-    {failed_sub_questions}
-    \n ------- \n
-
-    And here are the entities, relationships and terms extracted from the context:
-    \n ------- \n
-    {entity_term_extraction_str}
-    \n ------- \n
-
-   Please generate the list of good, fully contextualized sub-questions that would help to address the
-   main question. Again, please find questions that are NOT overlapping too much with the already answered
-   sub-questions or those that already were suggested and failed.
-   In other words - what can we try in addition to what has been tried so far?
-
-   Generate the list of json dictionaries with the following format:
-
-   {{"sub_questions": [{{"sub_question": <sub-question>,
-        "search_term": <rewrite the sub-question using as a search phrase for the document store>}},
-        ...]}} """
-
-DECOMPOSE_PROMPT = """ \n
-    For an initial user question, please generate at 5-10 individual sub-questions whose answers would help
-    \n to answer the initial question. The individual questions should be answerable by a good RAG system.
-    So a good idea would be to \n use the sub-questions to resolve ambiguities and/or to separate the
-    question for different entities that may be involved in the original question.
-
-    In order to arrive at meaningful sub-questions, please also consider the context retrieved from the
-    document store, expressed as entities, relationships and terms. You can also think about the types
-    mentioned in brackets
-
-    Guidelines:
-    - The sub-questions should be specific to the question and provide richer context for the question,
-    and or resolve ambiguities
-    - Each sub-question - when answered - should be relevant for the answer to the original question
-    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
-    other complications that may require extra context.
-    - The sub-questions MUST have the full context of the original question so that it can be executed by
-    a RAG system independently without the original question available
-      (Example:
-        - initial question: "What is the capital of France?"
-        - bad sub-question: "What is the name of the river there?"
-        - good sub-question: "What is the name of the river that flows through Paris?"
-    - For each sub-question, please provide a short explanation for why it is a good sub-question. So
-    generate a list of dictionaries with the following format:
-      [{{"sub_question": <sub-question>, "explanation": <explanation>, "search_term": <rewrite the
-      sub-question using as a search phrase for the document store>}}, ...]
-
-    \n\n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    And here are the entities, relationships and terms extracted from the context:
-    \n ------- \n
-    {entity_term_extraction_str}
-    \n ------- \n
-
-   Please generate the list of good, fully contextualized sub-questions that would help to address the
-   main question. Don't be too specific unless the original question is specific.
-   Please think through it step by step and then generate the list of json dictionaries with the following
-   format:
-   {{"sub_questions": [{{"sub_question": <sub-question>,
-        "explanation": <explanation>,
-        "search_term": <rewrite the sub-question using as a search phrase for the document store>}},
-        ...]}} """
-
-#### Consolidations
-COMBINED_CONTEXT = """-------
-    Below you will find useful information to answer the original question. First, you see a number of
-    sub-questions with their answers. This information should be considered to be more focussed and
-    somewhat more specific to the original question as it tries to contextualized facts.
-    After that will see the documents that were considered to be relevant to answer the original question.
-
-    Here are the sub-questions and their answers:
-    \n\n {deep_answer_context} \n\n
-    \n\n Here are the documents that were considered to be relevant to answer the original question:
-    \n\n {formated_docs} \n\n
-    ----------------
-    """
-
-SUB_QUESTION_EXPLANATION_RANKER_PROMPT = """-------
-    Below you will find a question that we ultimately want to answer (the original question) and a list of
-    motivations in arbitrary order for generated sub-questions that are supposed to help us answering the
-    original question. The motivations are formatted as <motivation number>:  <motivation explanation>.
-    (Again, the numbering is arbitrary and does not necessarily mean that 1 is the most relevant
-    motivation and 2 is less relevant.)
-
-    Please rank the motivations in order of relevance for answering the original question. Also, try to
-    ensure that the top questions do not duplicate too much, i.e. that they are not too similar.
-    Ultimately, create a list with the motivation numbers where the number of the most relevant
-    motivations comes first.
-
-    Here is the original question:
-    \n\n {original_question} \n\n
-    \n\n Here is the list of sub-question motivations:
-    \n\n {sub_question_explanations} \n\n
-    ----------------
-
-    Please think step by step and then generate the ranked list of motivations.
-
-    Please format your answer as a json object in the following format:
-    {{"reasonning": <explain your reasoning for the ranking>,
-      "ranked_motivations": <ranked list of motivation numbers>}}
-    """
-
-
-INITIAL_DECOMPOSITION_PROMPT = """ \n
-    Please decompose an initial user question into 2 or 3 appropriate sub-questions that help to
-    answer the original question. The purpose for this decomposition is to isolate individulal entities
-    (i.e., 'compare sales of company A and company B' -> 'what are sales for company A' + 'what are sales
-    for company B'), split ambiguous terms (i.e., 'what is our success with company A' -> 'what are our
-    sales with company A' + 'what is our market share with company A' + 'is company A a reference customer
-    for us'), etc. Each sub-question should be realistically be answerable by a good RAG system. \n
-
-    For each sub-question, please also create one search term that can be used to retrieve relevant
-    documents from a document store.
-
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Please formulate your answer as a list of json objects with the following format:
-
-   [{{"sub_question": <sub-question>, "search_term": <search term>}}, ...]
-
-    Answer:
-    """
-
-INITIAL_RAG_PROMPT = """ \n
-    You are an assistant for question-answering tasks. Use the information provided below - and only the
-    provided information - to answer the provided question.
-
-    The information provided below consists of:
-     1) a number of answered sub-questions - these are very important(!) and definitely should be
-     considered to answer the question.
-     2) a number of documents that were also deemed relevant for the question.
-
-    If you don't know the answer or if the provided information is empty or insufficient, just say
-    "I don't know". Do not use your internal knowledge!
-
-    Again, only use the provided informationand do not use your internal knowledge! It is a matter of life
-    and death that you do NOT use your internal knowledge, just the provided information!
-
-    Try to keep your answer concise.
-
-    And here is the question and the provided information:
-    \n
-    \nQuestion:\n {question}
-
-    \nAnswered Sub-questions:\n {answered_sub_questions}
-
-    \nContext:\n {context} \n\n
-    \n\n
-
-    Answer:"""
-
-ENTITY_TERM_PROMPT = """ \n
-    Based on the original question and the context retieved from a dataset, please generate a list of
-    entities (e.g. companies, organizations, industries, products, locations, etc.), terms and concepts
-    (e.g. sales, revenue, etc.) that are relevant for the question, plus their relations to each other.
-
-    \n\n
-    Here is the original question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-   And here is the context retrieved:
-    \n ------- \n
-    {context}
-    \n ------- \n
-
-    Please format your answer as a json object in the following format:
-
-    {{"retrieved_entities_relationships": {{
-        "entities": [{{
-            "entity_name": <assign a name for the entity>,
-            "entity_type": <specify a short type name for the entity, such as 'company', 'location',...>
-        }}],
-        "relationships": [{{
-            "name": <assign a name for the relationship>,
-            "type": <specify a short type name for the relationship, such as 'sales_to', 'is_location_of',...>,
-            "entities": [<related entity name 1>, <related entity name 2>]
-        }}],
-        "terms": [{{
-            "term_name": <assign a name for the term>,
-            "term_type": <specify a short type name for the term, such as 'revenue', 'market_share',...>,
-            "similar_to": <list terms that are similar to this term>
-        }}]
-    }}
-    }}
-   """

backend/danswer/agent_search/shared_graph_utils/utils.py

@@ -1,101 +0,0 @@
-import ast
-import json
-import re
-from collections.abc import Sequence
-from datetime import datetime
-from datetime import timedelta
-from typing import Any
-
-from danswer.context.search.models import InferenceSection
-
-
-def normalize_whitespace(text: str) -> str:
-    """Normalize whitespace in text to single spaces and strip leading/trailing whitespace."""
-    import re
-
-    return re.sub(r"\s+", " ", text.strip())
-
-
-# Post-processing
-def format_docs(docs: Sequence[InferenceSection]) -> str:
-    return "\n\n".join(doc.combined_content for doc in docs)
-
-
-def clean_and_parse_list_string(json_string: str) -> list[dict]:
-    # Remove any prefixes/labels before the actual JSON content
-    json_string = re.sub(r"^.*?(?=\[)", "", json_string, flags=re.DOTALL)
-
-    # Remove markdown code block markers and any newline prefixes
-    cleaned_string = re.sub(r"```json\n|\n```", "", json_string)
-    cleaned_string = cleaned_string.replace("\\n", " ").replace("\n", " ")
-    cleaned_string = " ".join(cleaned_string.split())
-
-    # Try parsing with json.loads first, fall back to ast.literal_eval
-    try:
-        return json.loads(cleaned_string)
-    except json.JSONDecodeError:
-        try:
-            return ast.literal_eval(cleaned_string)
-        except (ValueError, SyntaxError) as e:
-            raise ValueError(f"Failed to parse JSON string: {cleaned_string}") from e
-
-
-def clean_and_parse_json_string(json_string: str) -> dict[str, Any]:
-    # Remove markdown code block markers and any newline prefixes
-    cleaned_string = re.sub(r"```json\n|\n```", "", json_string)
-    cleaned_string = cleaned_string.replace("\\n", " ").replace("\n", " ")
-    cleaned_string = " ".join(cleaned_string.split())
-    # Parse the cleaned string into a Python dictionary
-    return json.loads(cleaned_string)
-
-
-def format_entity_term_extraction(entity_term_extraction_dict: dict[str, Any]) -> str:
-    entities = entity_term_extraction_dict["entities"]
-    terms = entity_term_extraction_dict["terms"]
-    relationships = entity_term_extraction_dict["relationships"]
-
-    entity_strs = ["\nEntities:\n"]
-    for entity in entities:
-        entity_str = f"{entity['entity_name']} ({entity['entity_type']})"
-        entity_strs.append(entity_str)
-
-    entity_str = "\n - ".join(entity_strs)
-
-    relationship_strs = ["\n\nRelationships:\n"]
-    for relationship in relationships:
-        relationship_str = f"{relationship['name']} ({relationship['type']}): {relationship['entities']}"
-        relationship_strs.append(relationship_str)
-
-    relationship_str = "\n - ".join(relationship_strs)
-
-    term_strs = ["\n\nTerms:\n"]
-    for term in terms:
-        term_str = f"{term['term_name']} ({term['term_type']}): similar to {term['similar_to']}"
-        term_strs.append(term_str)
-
-    term_str = "\n - ".join(term_strs)
-
-    return "\n".join(entity_strs + relationship_strs + term_strs)
-
-
-def _format_time_delta(time: timedelta) -> str:
-    seconds_from_start = f"{((time).seconds):03d}"
-    microseconds_from_start = f"{((time).microseconds):06d}"
-    return f"{seconds_from_start}.{microseconds_from_start}"
-
-
-def generate_log_message(
-    message: str,
-    node_start_time: datetime,
-    graph_start_time: datetime | None = None,
-) -> str:
-    current_time = datetime.now()
-
-    if graph_start_time is not None:
-        graph_time_str = _format_time_delta(current_time - graph_start_time)
-    else:
-        graph_time_str = "N/A"
-
-    node_time_str = _format_time_delta(current_time - node_start_time)
-
-    return f"{graph_time_str} ({node_time_str} s): {message}"

backend/danswer/auth/api_key.py

@@ -1,3 +1,4 @@
+import hashlib
 import secrets
 import uuid
 from urllib.parse import quote
@@ -18,7 +19,8 @@ _API_KEY_HEADER_NAME = "Authorization"
 # organizations like the Internet Engineering Task Force (IETF).
 _API_KEY_HEADER_ALTERNATIVE_NAME = "X-Danswer-Authorization"
 _BEARER_PREFIX = "Bearer "
-_API_KEY_PREFIX = "dn_"
+_API_KEY_PREFIX = "on_"
+_DEPRECATED_API_KEY_PREFIX = "dn_"
 _API_KEY_LEN = 192
 
 
@@ -52,7 +54,9 @@ def extract_tenant_from_api_key_header(request: Request) -> str | None:
 
     api_key = raw_api_key_header[len(_BEARER_PREFIX) :].strip()
 
-    if not api_key.startswith(_API_KEY_PREFIX):
+    if not api_key.startswith(_API_KEY_PREFIX) and not api_key.startswith(
+        _DEPRECATED_API_KEY_PREFIX
+    ):
         return None
 
     parts = api_key[len(_API_KEY_PREFIX) :].split(".", 1)
@@ -63,10 +67,19 @@ def extract_tenant_from_api_key_header(request: Request) -> str | None:
     return unquote(tenant_id) if tenant_id else None
 
 
+def _deprecated_hash_api_key(api_key: str) -> str:
+    return sha256_crypt.hash(api_key, salt="", rounds=API_KEY_HASH_ROUNDS)
+
+
 def hash_api_key(api_key: str) -> str:
     # NOTE: no salt is needed, as the API key is randomly generated
     # and overlaps are impossible
-    return sha256_crypt.hash(api_key, salt="", rounds=API_KEY_HASH_ROUNDS)
+    if api_key.startswith(_API_KEY_PREFIX):
+        return hashlib.sha256(api_key.encode("utf-8")).hexdigest()
+    elif api_key.startswith(_DEPRECATED_API_KEY_PREFIX):
+        return _deprecated_hash_api_key(api_key)
+    else:
+        raise ValueError(f"Invalid API key prefix: {api_key[:3]}")
 
 
 def build_displayable_api_key(api_key: str) -> str:

backend/danswer/auth/invited_users.py

@@ -9,7 +9,6 @@ from danswer.utils.special_types import JSON_ro
 def get_invited_users() -> list[str]:
     try:
         store = get_kv_store()
-
         return cast(list, store.load(KV_USER_STORE_KEY))
     except KvKeyNotFoundError:
         return list()

backend/danswer/auth/users.py

@@ -58,7 +58,6 @@ from danswer.auth.schemas import UserRole
 from danswer.auth.schemas import UserUpdate
 from danswer.configs.app_configs import AUTH_TYPE
 from danswer.configs.app_configs import DISABLE_AUTH
-from danswer.configs.app_configs import DISABLE_VERIFICATION
 from danswer.configs.app_configs import EMAIL_FROM
 from danswer.configs.app_configs import REQUIRE_EMAIL_VERIFICATION
 from danswer.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS
@@ -132,11 +131,12 @@ def get_display_email(email: str | None, space_less: bool = False) -> str:
 
 
 def user_needs_to_be_verified() -> bool:
-    # all other auth types besides basic should require users to be
-    # verified
-    return not DISABLE_VERIFICATION and (
-        AUTH_TYPE != AuthType.BASIC or REQUIRE_EMAIL_VERIFICATION
-    )
+    if AUTH_TYPE == AuthType.BASIC:
+        return REQUIRE_EMAIL_VERIFICATION
+
+    # For other auth types, if the user is authenticated it's assumed that
+    # the user is already verified via the external IDP
+    return False
 
 
 def verify_email_is_invited(email: str) -> None:

backend/danswer/background/celery/tasks/doc_permission_syncing/tasks.py

@@ -219,7 +219,7 @@ def connector_permission_sync_generator_task(
 
     r = get_redis_client(tenant_id=tenant_id)
 
-    lock = r.lock(
+    lock: RedisLock = r.lock(
         DanswerRedisLocks.CONNECTOR_DOC_PERMISSIONS_SYNC_LOCK_PREFIX
         + f"_{redis_connector.id}",
         timeout=CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT,

backend/danswer/background/celery/tasks/indexing/tasks.py

@@ -598,7 +598,7 @@ def connector_indexing_proxy_task(
                         db_session,
                         "Connector termination signal detected",
                     )
-            finally:
+            except Exception:
                 # if the DB exceptions, we'll just get an unfriendly failure message
                 # in the UI instead of the cancellation message
                 logger.exception(
@@ -640,12 +640,16 @@ def connector_indexing_proxy_task(
                 continue
 
         if job.status == "error":
+            exit_code: int | None = None
+            if job.process:
+                exit_code = job.process.exitcode
             task_logger.error(
                 "Indexing watchdog - spawned task exceptioned: "
                 f"attempt={index_attempt_id} "
                 f"tenant={tenant_id} "
                 f"cc_pair={cc_pair_id} "
                 f"search_settings={search_settings_id} "
+                f"exit_code={exit_code} "
                 f"error={job.exception()}"
             )
 

backend/danswer/background/celery/tasks/vespa/tasks.py

@@ -64,6 +64,7 @@ from danswer.redis.redis_connector_doc_perm_sync import (
 )
 from danswer.redis.redis_connector_index import RedisConnectorIndex
 from danswer.redis.redis_connector_prune import RedisConnectorPrune
+from danswer.redis.redis_connector_utils import RedisConnectorUtils
 from danswer.redis.redis_document_set import RedisDocumentSet
 from danswer.redis.redis_pool import get_redis_client
 from danswer.redis.redis_usergroup import RedisUserGroup
@@ -402,7 +403,7 @@ def monitor_connector_deletion_taskset(
     tenant_id: str | None, key_bytes: bytes, r: Redis
 ) -> None:
     fence_key = key_bytes.decode("utf-8")
-    cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)
+    cc_pair_id_str = RedisConnectorUtils.get_id_from_fence_key(fence_key)
     if cc_pair_id_str is None:
         task_logger.warning(f"could not parse cc_pair_id from {fence_key}")
         return
@@ -528,7 +529,7 @@ def monitor_ccpair_pruning_taskset(
     tenant_id: str | None, key_bytes: bytes, r: Redis, db_session: Session
 ) -> None:
     fence_key = key_bytes.decode("utf-8")
-    cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)
+    cc_pair_id_str = RedisConnectorUtils.get_id_from_fence_key(fence_key)
     if cc_pair_id_str is None:
         task_logger.warning(
             f"monitor_ccpair_pruning_taskset: could not parse cc_pair_id from {fence_key}"
@@ -566,7 +567,7 @@ def monitor_ccpair_permissions_taskset(
     tenant_id: str | None, key_bytes: bytes, r: Redis, db_session: Session
 ) -> None:
     fence_key = key_bytes.decode("utf-8")
-    cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)
+    cc_pair_id_str = RedisConnectorUtils.get_id_from_fence_key(fence_key)
     if cc_pair_id_str is None:
         task_logger.warning(
             f"monitor_ccpair_permissions_taskset: could not parse cc_pair_id from {fence_key}"
@@ -604,25 +605,15 @@ def monitor_ccpair_permissions_taskset(
 def monitor_ccpair_indexing_taskset(
     tenant_id: str | None, key_bytes: bytes, r: Redis, db_session: Session
 ) -> None:
-    # if the fence doesn't exist, there's nothing to do
-    fence_key = key_bytes.decode("utf-8")
-    composite_id = RedisConnector.get_id_from_fence_key(fence_key)
-    if composite_id is None:
+    redis_ids = RedisConnectorIndex.parse_key(key_bytes)
+    if not redis_ids:
         task_logger.warning(
-            f"monitor_ccpair_indexing_taskset: could not parse composite_id from {fence_key}"
+            f"monitor_ccpair_indexing_taskset: could not parse composite_id from {key_bytes!r}"
         )
         return
 
-    # parse out metadata and initialize the helper class with it
-    parts = composite_id.split("/")
-    if len(parts) != 2:
-        return
-
-    cc_pair_id = int(parts[0])
-    search_settings_id = int(parts[1])
-
-    redis_connector = RedisConnector(tenant_id, cc_pair_id)
-    redis_connector_index = redis_connector.new_index(search_settings_id)
+    redis_connector = RedisConnector(tenant_id, redis_ids.cc_pair_id)
+    redis_connector_index = redis_connector.new_index(redis_ids.search_settings_id)
     if not redis_connector_index.fenced:
         return
 
@@ -635,8 +626,8 @@ def monitor_ccpair_indexing_taskset(
     progress = redis_connector_index.get_progress()
     if progress is not None:
         task_logger.info(
-            f"Connector indexing progress: cc_pair={cc_pair_id} "
-            f"search_settings={search_settings_id} "
+            f"Connector indexing progress: cc_pair={redis_ids.cc_pair_id} "
+            f"search_settings={redis_ids.search_settings_id} "
             f"progress={progress} "
             f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
         )
@@ -671,8 +662,8 @@ def monitor_ccpair_indexing_taskset(
                     f"Connector indexing aborted or exceptioned: "
                     f"attempt={payload.index_attempt_id} "
                     f"celery_task={payload.celery_task_id} "
-                    f"cc_pair={cc_pair_id} "
-                    f"search_settings={search_settings_id} "
+                    f"cc_pair={redis_ids.cc_pair_id} "
+                    f"search_settings={redis_ids.search_settings_id} "
                     f"elapsed_submitted={elapsed_submitted.total_seconds():.2f} "
                     f"result.state={task_state} "
                     f"result.result={task_result} "
@@ -680,17 +671,28 @@ def monitor_ccpair_indexing_taskset(
                 )
                 task_logger.warning(msg)
 
-                index_attempt = get_index_attempt(db_session, payload.index_attempt_id)
-                if index_attempt:
-                    if (
-                        index_attempt.status != IndexingStatus.CANCELED
-                        and index_attempt.status != IndexingStatus.FAILED
-                    ):
-                        mark_attempt_failed(
-                            index_attempt_id=payload.index_attempt_id,
-                            db_session=db_session,
-                            failure_reason=msg,
-                        )
+                try:
+                    index_attempt = get_index_attempt(
+                        db_session, payload.index_attempt_id
+                    )
+                    if index_attempt:
+                        if (
+                            index_attempt.status != IndexingStatus.CANCELED
+                            and index_attempt.status != IndexingStatus.FAILED
+                        ):
+                            mark_attempt_failed(
+                                index_attempt_id=payload.index_attempt_id,
+                                db_session=db_session,
+                                failure_reason=msg,
+                            )
+                except Exception:
+                    task_logger.exception(
+                        "monitor_ccpair_indexing_taskset - transient exception marking index attempt as failed: "
+                        f"attempt={payload.index_attempt_id} "
+                        f"tenant={tenant_id} "
+                        f"cc_pair={redis_ids.cc_pair_id} "
+                        f"search_settings={redis_ids.search_settings_id}"
+                    )
 
                 redis_connector_index.reset()
         return
@@ -698,8 +700,8 @@ def monitor_ccpair_indexing_taskset(
     status_enum = HTTPStatus(status_int)
 
     task_logger.info(
-        f"Connector indexing finished: cc_pair={cc_pair_id} "
-        f"search_settings={search_settings_id} "
+        f"Connector indexing finished: cc_pair={redis_ids.cc_pair_id} "
+        f"search_settings={redis_ids.search_settings_id} "
         f"progress={progress} "
         f"status={status_enum.name} "
         f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"

backend/danswer/background/indexing/job_client.py

@@ -82,7 +82,7 @@ class SimpleJob:
             return "running"
         elif self.process.exitcode is None:
             return "cancelled"
-        elif self.process.exitcode > 0:
+        elif self.process.exitcode != 0:
             return "error"
         else:
             return "finished"
@@ -123,7 +123,8 @@ class SimpleJobClient:
         self._cleanup_completed_jobs()
         if len(self.jobs) >= self.n_workers:
             logger.debug(
-                f"No available workers to run job. Currently running '{len(self.jobs)}' jobs, with a limit of '{self.n_workers}'."
+                f"No available workers to run job. "
+                f"Currently running '{len(self.jobs)}' jobs, with a limit of '{self.n_workers}'."
             )
             return None
 

backend/danswer/chat/answer.py

@@ -6,27 +6,27 @@ from langchain.schema.messages import BaseMessage
 from langchain_core.messages import AIMessageChunk
 from langchain_core.messages import ToolCall
 
+from danswer.chat.llm_response_handler import LLMResponseHandlerManager
 from danswer.chat.models import AnswerQuestionPossibleReturn
+from danswer.chat.models import AnswerStyleConfig
 from danswer.chat.models import CitationInfo
 from danswer.chat.models import DanswerAnswerPiece
-from danswer.file_store.utils import InMemoryChatFile
-from danswer.llm.answering.llm_response_handler import LLMCall
-from danswer.llm.answering.llm_response_handler import LLMResponseHandlerManager
-from danswer.llm.answering.models import AnswerStyleConfig
-from danswer.llm.answering.models import PreviousMessage
-from danswer.llm.answering.models import PromptConfig
-from danswer.llm.answering.prompts.build import AnswerPromptBuilder
-from danswer.llm.answering.prompts.build import default_build_system_message
-from danswer.llm.answering.prompts.build import default_build_user_message
-from danswer.llm.answering.stream_processing.answer_response_handler import (
+from danswer.chat.models import PromptConfig
+from danswer.chat.prompt_builder.build import AnswerPromptBuilder
+from danswer.chat.prompt_builder.build import default_build_system_message
+from danswer.chat.prompt_builder.build import default_build_user_message
+from danswer.chat.prompt_builder.build import LLMCall
+from danswer.chat.stream_processing.answer_response_handler import (
     CitationResponseHandler,
 )
-from danswer.llm.answering.stream_processing.answer_response_handler import (
+from danswer.chat.stream_processing.answer_response_handler import (
     DummyAnswerResponseHandler,
 )
-from danswer.llm.answering.stream_processing.utils import map_document_id_order
-from danswer.llm.answering.tool.tool_response_handler import ToolResponseHandler
+from danswer.chat.stream_processing.utils import map_document_id_order
+from danswer.chat.tool_handling.tool_response_handler import ToolResponseHandler
+from danswer.file_store.utils import InMemoryChatFile
 from danswer.llm.interfaces import LLM
+from danswer.llm.models import PreviousMessage
 from danswer.natural_language_processing.utils import get_tokenizer
 from danswer.tools.force import ForceUseTool
 from danswer.tools.models import ToolResponse
@@ -206,7 +206,9 @@ class Answer:
         # + figure out what the next LLM call should be
         tool_call_handler = ToolResponseHandler(current_llm_call.tools)
 
-        search_result = SearchTool.get_search_result(current_llm_call) or []
+        search_result, displayed_search_results_map = SearchTool.get_search_result(
+            current_llm_call
+        ) or ([], {})
 
         # Quotes are no longer supported
         # answer_handler: AnswerResponseHandler
@@ -224,6 +226,7 @@ class Answer:
         answer_handler = CitationResponseHandler(
             context_docs=search_result,
             doc_id_to_rank_map=map_document_id_order(search_result),
+            display_doc_order_dict=displayed_search_results_map,
         )
 
         response_handler_manager = LLMResponseHandlerManager(

backend/danswer/chat/chat_utils.py

@@ -26,7 +26,7 @@ from danswer.db.models import Prompt
 from danswer.db.models import Tool
 from danswer.db.models import User
 from danswer.db.persona import get_prompts_by_ids
-from danswer.llm.answering.models import PreviousMessage
+from danswer.llm.models import PreviousMessage
 from danswer.natural_language_processing.utils import BaseTokenizer
 from danswer.server.query_and_chat.models import CreateChatMessageRequest
 from danswer.tools.tool_implementations.custom.custom_tool import (

backend/danswer/chat/llm_response_handler.py

@@ -1,58 +1,22 @@
 from collections.abc import Callable
 from collections.abc import Generator
 from collections.abc import Iterator
-from typing import TYPE_CHECKING
 
 from langchain_core.messages import BaseMessage
-from pydantic.v1 import BaseModel as BaseModel__v1
 
-from danswer.chat.models import CitationInfo
-from danswer.chat.models import DanswerAnswerPiece
+from danswer.chat.models import ResponsePart
 from danswer.chat.models import StreamStopInfo
 from danswer.chat.models import StreamStopReason
-from danswer.file_store.models import InMemoryChatFile
-from danswer.llm.answering.prompts.build import AnswerPromptBuilder
-from danswer.tools.force import ForceUseTool
-from danswer.tools.models import ToolCallFinalResult
-from danswer.tools.models import ToolCallKickoff
-from danswer.tools.models import ToolResponse
-from danswer.tools.tool import Tool
-
-
-if TYPE_CHECKING:
-    from danswer.llm.answering.stream_processing.answer_response_handler import (
-        AnswerResponseHandler,
-    )
-    from danswer.llm.answering.tool.tool_response_handler import ToolResponseHandler
-
-
-ResponsePart = (
-    DanswerAnswerPiece
-    | CitationInfo
-    | ToolCallKickoff
-    | ToolResponse
-    | ToolCallFinalResult
-    | StreamStopInfo
-)
-
-
-class LLMCall(BaseModel__v1):
-    prompt_builder: AnswerPromptBuilder
-    tools: list[Tool]
-    force_use_tool: ForceUseTool
-    files: list[InMemoryChatFile]
-    tool_call_info: list[ToolCallKickoff | ToolResponse | ToolCallFinalResult]
-    using_tool_calling_llm: bool
-
-    class Config:
-        arbitrary_types_allowed = True
+from danswer.chat.prompt_builder.build import LLMCall
+from danswer.chat.stream_processing.answer_response_handler import AnswerResponseHandler
+from danswer.chat.tool_handling.tool_response_handler import ToolResponseHandler
 
 
 class LLMResponseHandlerManager:
     def __init__(
         self,
-        tool_handler: "ToolResponseHandler",
-        answer_handler: "AnswerResponseHandler",
+        tool_handler: ToolResponseHandler,
+        answer_handler: AnswerResponseHandler,
         is_cancelled: Callable[[], bool],
     ):
         self.tool_handler = tool_handler

backend/danswer/chat/models.py

@@ -1,10 +1,14 @@
+from collections.abc import Callable
 from collections.abc import Iterator
 from datetime import datetime
 from enum import Enum
 from typing import Any
+from typing import TYPE_CHECKING
 
 from pydantic import BaseModel
+from pydantic import ConfigDict
 from pydantic import Field
+from pydantic import model_validator
 
 from danswer.configs.constants import DocumentSource
 from danswer.configs.constants import MessageType
@@ -12,8 +16,15 @@ from danswer.context.search.enums import QueryFlow
 from danswer.context.search.enums import RecencyBiasSetting
 from danswer.context.search.enums import SearchType
 from danswer.context.search.models import RetrievalDocs
+from danswer.llm.override_models import PromptOverride
+from danswer.tools.models import ToolCallFinalResult
+from danswer.tools.models import ToolCallKickoff
+from danswer.tools.models import ToolResponse
 from danswer.tools.tool_implementations.custom.base_tool_types import ToolResultType
 
+if TYPE_CHECKING:
+    from danswer.db.models import Prompt
+
 
 class LlmDoc(BaseModel):
     """This contains the minimal set information for the LLM portion including citations"""
@@ -210,3 +221,109 @@ AnswerQuestionStreamReturn = Iterator[AnswerQuestionPossibleReturn]
 class LLMMetricsContainer(BaseModel):
     prompt_tokens: int
     response_tokens: int
+
+
+StreamProcessor = Callable[[Iterator[str]], AnswerQuestionStreamReturn]
+
+
+class DocumentPruningConfig(BaseModel):
+    max_chunks: int | None = None
+    max_window_percentage: float | None = None
+    max_tokens: int | None = None
+    # different pruning behavior is expected when the
+    # user manually selects documents they want to chat with
+    # e.g. we don't want to truncate each document to be no more
+    # than one chunk long
+    is_manually_selected_docs: bool = False
+    # If user specifies to include additional context Chunks for each match, then different pruning
+    # is used. As many Sections as possible are included, and the last Section is truncated
+    # If this is false, all of the Sections are truncated if they are longer than the expected Chunk size.
+    # Sections are often expected to be longer than the maximum Chunk size but Chunks should not be.
+    use_sections: bool = True
+    # If using tools, then we need to consider the tool length
+    tool_num_tokens: int = 0
+    # If using a tool message to represent the docs, then we have to JSON serialize
+    # the document content, which adds to the token count.
+    using_tool_message: bool = False
+
+
+class ContextualPruningConfig(DocumentPruningConfig):
+    num_chunk_multiple: int
+
+    @classmethod
+    def from_doc_pruning_config(
+        cls, num_chunk_multiple: int, doc_pruning_config: DocumentPruningConfig
+    ) -> "ContextualPruningConfig":
+        return cls(num_chunk_multiple=num_chunk_multiple, **doc_pruning_config.dict())
+
+
+class CitationConfig(BaseModel):
+    all_docs_useful: bool = False
+
+
+class QuotesConfig(BaseModel):
+    pass
+
+
+class AnswerStyleConfig(BaseModel):
+    citation_config: CitationConfig | None = None
+    quotes_config: QuotesConfig | None = None
+    document_pruning_config: DocumentPruningConfig = Field(
+        default_factory=DocumentPruningConfig
+    )
+    # forces the LLM to return a structured response, see
+    # https://platform.openai.com/docs/guides/structured-outputs/introduction
+    # right now, only used by the simple chat API
+    structured_response_format: dict | None = None
+
+    @model_validator(mode="after")
+    def check_quotes_and_citation(self) -> "AnswerStyleConfig":
+        if self.citation_config is None and self.quotes_config is None:
+            raise ValueError(
+                "One of `citation_config` or `quotes_config` must be provided"
+            )
+
+        if self.citation_config is not None and self.quotes_config is not None:
+            raise ValueError(
+                "Only one of `citation_config` or `quotes_config` must be provided"
+            )
+
+        return self
+
+
+class PromptConfig(BaseModel):
+    """Final representation of the Prompt configuration passed
+    into the `Answer` object."""
+
+    system_prompt: str
+    task_prompt: str
+    datetime_aware: bool
+    include_citations: bool
+
+    @classmethod
+    def from_model(
+        cls, model: "Prompt", prompt_override: PromptOverride | None = None
+    ) -> "PromptConfig":
+        override_system_prompt = (
+            prompt_override.system_prompt if prompt_override else None
+        )
+        override_task_prompt = prompt_override.task_prompt if prompt_override else None
+
+        return cls(
+            system_prompt=override_system_prompt or model.system_prompt,
+            task_prompt=override_task_prompt or model.task_prompt,
+            datetime_aware=model.datetime_aware,
+            include_citations=model.include_citations,
+        )
+
+    model_config = ConfigDict(frozen=True)
+
+
+ResponsePart = (
+    DanswerAnswerPiece
+    | CitationInfo
+    | ToolCallKickoff
+    | ToolResponse
+    | ToolCallFinalResult
+    | StreamStopInfo
+)

backend/danswer/chat/process_message.py

@@ -6,19 +6,24 @@ from typing import cast
 
 from sqlalchemy.orm import Session
 
+from danswer.chat.answer import Answer
 from danswer.chat.chat_utils import create_chat_chain
 from danswer.chat.chat_utils import create_temporary_persona
 from danswer.chat.models import AllCitations
+from danswer.chat.models import AnswerStyleConfig
 from danswer.chat.models import ChatDanswerBotResponse
+from danswer.chat.models import CitationConfig
 from danswer.chat.models import CitationInfo
 from danswer.chat.models import CustomToolResponse
 from danswer.chat.models import DanswerAnswerPiece
 from danswer.chat.models import DanswerContexts
+from danswer.chat.models import DocumentPruningConfig
 from danswer.chat.models import FileChatDisplay
 from danswer.chat.models import FinalUsedContextDocsResponse
 from danswer.chat.models import LLMRelevanceFilterResponse
 from danswer.chat.models import MessageResponseIDInfo
 from danswer.chat.models import MessageSpecificCitations
+from danswer.chat.models import PromptConfig
 from danswer.chat.models import QADocsResponse
 from danswer.chat.models import StreamingError
 from danswer.chat.models import StreamStopInfo
@@ -57,16 +62,11 @@ from danswer.document_index.factory import get_default_document_index
 from danswer.file_store.models import ChatFileType
 from danswer.file_store.models import FileDescriptor
 from danswer.file_store.utils import load_all_chat_files
-from danswer.file_store.utils import save_files_from_urls
-from danswer.llm.answering.answer import Answer
-from danswer.llm.answering.models import AnswerStyleConfig
-from danswer.llm.answering.models import CitationConfig
-from danswer.llm.answering.models import DocumentPruningConfig
-from danswer.llm.answering.models import PreviousMessage
-from danswer.llm.answering.models import PromptConfig
+from danswer.file_store.utils import save_files
 from danswer.llm.exceptions import GenAIDisabledException
 from danswer.llm.factory import get_llms_for_persona
 from danswer.llm.factory import get_main_llm_from_tuple
+from danswer.llm.models import PreviousMessage
 from danswer.llm.utils import litellm_exception_to_error_msg
 from danswer.natural_language_processing.utils import get_tokenizer
 from danswer.server.query_and_chat.models import ChatMessageDetail
@@ -119,6 +119,7 @@ from danswer.utils.logger import setup_logger
 from danswer.utils.long_term_log import LongTermLogger
 from danswer.utils.timing import log_function_time
 from danswer.utils.timing import log_generator_function_time
+from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 
 
 logger = setup_logger()
@@ -302,6 +303,7 @@ def stream_chat_message_objects(
     3. [always] A set of streamed LLM tokens or an error anywhere along the line if something fails
     4. [always] Details on the final AI response message that is created
     """
+    tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
     use_existing_user_message = new_msg_req.use_existing_user_message
     existing_assistant_message_id = new_msg_req.existing_assistant_message_id
 
@@ -678,7 +680,8 @@ def stream_chat_message_objects(
 
         reference_db_search_docs = None
         qa_docs_response = None
-        ai_message_files = None  # any files to associate with the AI message e.g. dall-e generated images
+        # any files to associate with the AI message e.g. dall-e generated images
+        ai_message_files = []
         dropped_indices = None
         tool_result = None
 
@@ -733,8 +736,14 @@ def stream_chat_message_objects(
                         list[ImageGenerationResponse], packet.response
                     )
 
-                    file_ids = save_files_from_urls(
-                        [img.url for img in img_generation_response]
+                    file_ids = save_files(
+                        urls=[img.url for img in img_generation_response if img.url],
+                        base64_files=[
+                            img.image_data
+                            for img in img_generation_response
+                            if img.image_data
+                        ],
+                        tenant_id=tenant_id,
                     )
                     ai_message_files = [
                         FileDescriptor(id=str(file_id), type=ChatFileType.IMAGE)
@@ -760,15 +769,19 @@ def stream_chat_message_objects(
                         or custom_tool_response.response_type == "csv"
                     ):
                         file_ids = custom_tool_response.tool_result.file_ids
-                        ai_message_files = [
-                            FileDescriptor(
-                                id=str(file_id),
-                                type=ChatFileType.IMAGE
-                                if custom_tool_response.response_type == "image"
-                                else ChatFileType.CSV,
-                            )
-                            for file_id in file_ids
-                        ]
+                        ai_message_files.extend(
+                            [
+                                FileDescriptor(
+                                    id=str(file_id),
+                                    type=(
+                                        ChatFileType.IMAGE
+                                        if custom_tool_response.response_type == "image"
+                                        else ChatFileType.CSV
+                                    ),
+                                )
+                                for file_id in file_ids
+                            ]
+                        )
                         yield FileChatDisplay(
                             file_ids=[str(file_id) for file_id in file_ids]
                         )
@@ -818,7 +831,8 @@ def stream_chat_message_objects(
                 citations_list=answer.citations,
                 db_docs=reference_db_search_docs,
             )
-            yield AllCitations(citations=answer.citations)
+            if not answer.is_cancelled():
+                yield AllCitations(citations=answer.citations)
 
         # Saving Gen AI answer and responding with message info
         tool_name_to_tool_id: dict[str, int] = {}

backend/danswer/chat/prompt_builder/build.py

@@ -4,20 +4,26 @@ from typing import cast
 from langchain_core.messages import BaseMessage
 from langchain_core.messages import HumanMessage
 from langchain_core.messages import SystemMessage
+from pydantic.v1 import BaseModel as BaseModel__v1
 
+from danswer.chat.models import PromptConfig
+from danswer.chat.prompt_builder.citations_prompt import compute_max_llm_input_tokens
+from danswer.chat.prompt_builder.utils import translate_history_to_basemessages
 from danswer.file_store.models import InMemoryChatFile
-from danswer.llm.answering.models import PreviousMessage
-from danswer.llm.answering.models import PromptConfig
-from danswer.llm.answering.prompts.citations_prompt import compute_max_llm_input_tokens
 from danswer.llm.interfaces import LLMConfig
+from danswer.llm.models import PreviousMessage
 from danswer.llm.utils import build_content_with_imgs
 from danswer.llm.utils import check_message_tokens
 from danswer.llm.utils import message_to_prompt_and_imgs
-from danswer.llm.utils import translate_history_to_basemessages
 from danswer.natural_language_processing.utils import get_tokenizer
 from danswer.prompts.chat_prompts import CHAT_USER_CONTEXT_FREE_PROMPT
 from danswer.prompts.prompt_utils import add_date_time_to_prompt
 from danswer.prompts.prompt_utils import drop_messages_history_overflow
+from danswer.tools.force import ForceUseTool
+from danswer.tools.models import ToolCallFinalResult
+from danswer.tools.models import ToolCallKickoff
+from danswer.tools.models import ToolResponse
+from danswer.tools.tool import Tool
 
 
 def default_build_system_message(
@@ -139,3 +145,15 @@ class AnswerPromptBuilder:
         return drop_messages_history_overflow(
             final_messages_with_tokens, self.max_tokens
         )
+
+
+class LLMCall(BaseModel__v1):
+    prompt_builder: AnswerPromptBuilder
+    tools: list[Tool]
+    force_use_tool: ForceUseTool
+    files: list[InMemoryChatFile]
+    tool_call_info: list[ToolCallKickoff | ToolResponse | ToolCallFinalResult]
+    using_tool_calling_llm: bool
+
+    class Config:
+        arbitrary_types_allowed = True

backend/danswer/chat/prompt_builder/citations_prompt.py

@@ -2,12 +2,12 @@ from langchain.schema.messages import HumanMessage
 from langchain.schema.messages import SystemMessage
 
 from danswer.chat.models import LlmDoc
+from danswer.chat.models import PromptConfig
 from danswer.configs.model_configs import GEN_AI_SINGLE_USER_MESSAGE_EXPECTED_MAX_TOKENS
 from danswer.context.search.models import InferenceChunk
 from danswer.db.models import Persona
 from danswer.db.persona import get_default_prompt__read_only
 from danswer.db.search_settings import get_multilingual_expansion
-from danswer.llm.answering.models import PromptConfig
 from danswer.llm.factory import get_llms_for_persona
 from danswer.llm.factory import get_main_llm_from_tuple
 from danswer.llm.interfaces import LLMConfig

backend/danswer/chat/prompt_builder/quotes_prompt.py

@@ -1,10 +1,10 @@
 from langchain.schema.messages import HumanMessage
 
 from danswer.chat.models import LlmDoc
+from danswer.chat.models import PromptConfig
 from danswer.configs.chat_configs import LANGUAGE_HINT
 from danswer.context.search.models import InferenceChunk
 from danswer.db.search_settings import get_multilingual_expansion
-from danswer.llm.answering.models import PromptConfig
 from danswer.llm.utils import message_to_prompt_and_imgs
 from danswer.prompts.direct_qa_prompts import CONTEXT_BLOCK
 from danswer.prompts.direct_qa_prompts import HISTORY_BLOCK

backend/danswer/chat/prompt_builder/utils.py

@@ -0,0 +1,62 @@
+from langchain.schema.messages import AIMessage
+from langchain.schema.messages import BaseMessage
+from langchain.schema.messages import HumanMessage
+
+from danswer.configs.constants import MessageType
+from danswer.db.models import ChatMessage
+from danswer.file_store.models import InMemoryChatFile
+from danswer.llm.models import PreviousMessage
+from danswer.llm.utils import build_content_with_imgs
+from danswer.prompts.direct_qa_prompts import PARAMATERIZED_PROMPT
+from danswer.prompts.direct_qa_prompts import PARAMATERIZED_PROMPT_WITHOUT_CONTEXT
+
+
+def build_dummy_prompt(
+    system_prompt: str, task_prompt: str, retrieval_disabled: bool
+) -> str:
+    if retrieval_disabled:
+        return PARAMATERIZED_PROMPT_WITHOUT_CONTEXT.format(
+            user_query="<USER_QUERY>",
+            system_prompt=system_prompt,
+            task_prompt=task_prompt,
+        ).strip()
+
+    return PARAMATERIZED_PROMPT.format(
+        context_docs_str="<CONTEXT_DOCS>",
+        user_query="<USER_QUERY>",
+        system_prompt=system_prompt,
+        task_prompt=task_prompt,
+    ).strip()
+
+
+def translate_danswer_msg_to_langchain(
+    msg: ChatMessage | PreviousMessage,
+) -> BaseMessage:
+    files: list[InMemoryChatFile] = []
+
+    # If the message is a `ChatMessage`, it doesn't have the downloaded files
+    # attached. Just ignore them for now.
+    if not isinstance(msg, ChatMessage):
+        files = msg.files
+    content = build_content_with_imgs(msg.message, files, message_type=msg.message_type)
+
+    if msg.message_type == MessageType.SYSTEM:
+        raise ValueError("System messages are not currently part of history")
+    if msg.message_type == MessageType.ASSISTANT:
+        return AIMessage(content=content)
+    if msg.message_type == MessageType.USER:
+        return HumanMessage(content=content)
+
+    raise ValueError(f"New message type {msg.message_type} not handled")
+
+
+def translate_history_to_basemessages(
+    history: list[ChatMessage] | list["PreviousMessage"],
+) -> tuple[list[BaseMessage], list[int]]:
+    history_basemessages = [
+        translate_danswer_msg_to_langchain(msg)
+        for msg in history
+        if msg.token_count != 0
+    ]
+    history_token_counts = [msg.token_count for msg in history if msg.token_count != 0]
+    return history_basemessages, history_token_counts

backend/danswer/chat/prune_and_merge.py

@@ -5,16 +5,16 @@ from typing import TypeVar
 
 from pydantic import BaseModel
 
+from danswer.chat.models import ContextualPruningConfig
 from danswer.chat.models import (
     LlmDoc,
 )
+from danswer.chat.models import PromptConfig
+from danswer.chat.prompt_builder.citations_prompt import compute_max_document_tokens
 from danswer.configs.constants import IGNORE_FOR_QA
 from danswer.configs.model_configs import DOC_EMBEDDING_CONTEXT_SIZE
 from danswer.context.search.models import InferenceChunk
 from danswer.context.search.models import InferenceSection
-from danswer.llm.answering.models import ContextualPruningConfig
-from danswer.llm.answering.models import PromptConfig
-from danswer.llm.answering.prompts.citations_prompt import compute_max_document_tokens
 from danswer.llm.interfaces import LLMConfig
 from danswer.natural_language_processing.utils import get_tokenizer
 from danswer.natural_language_processing.utils import tokenizer_trim_content

backend/danswer/chat/stream_processing/answer_response_handler.py

@@ -3,13 +3,11 @@ from collections.abc import Generator
 
 from langchain_core.messages import BaseMessage
 
+from danswer.chat.llm_response_handler import ResponsePart
 from danswer.chat.models import CitationInfo
 from danswer.chat.models import LlmDoc
-from danswer.llm.answering.llm_response_handler import ResponsePart
-from danswer.llm.answering.stream_processing.citation_processing import (
-    CitationProcessor,
-)
-from danswer.llm.answering.stream_processing.utils import DocumentIdOrderMapping
+from danswer.chat.stream_processing.citation_processing import CitationProcessor
+from danswer.chat.stream_processing.utils import DocumentIdOrderMapping
 from danswer.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -37,13 +35,18 @@ class DummyAnswerResponseHandler(AnswerResponseHandler):
 
 class CitationResponseHandler(AnswerResponseHandler):
     def __init__(
-        self, context_docs: list[LlmDoc], doc_id_to_rank_map: DocumentIdOrderMapping
+        self,
+        context_docs: list[LlmDoc],
+        doc_id_to_rank_map: DocumentIdOrderMapping,
+        display_doc_order_dict: dict[str, int],
     ):
         self.context_docs = context_docs
         self.doc_id_to_rank_map = doc_id_to_rank_map
+        self.display_doc_order_dict = display_doc_order_dict
         self.citation_processor = CitationProcessor(
             context_docs=self.context_docs,
             doc_id_to_rank_map=self.doc_id_to_rank_map,
+            display_doc_order_dict=self.display_doc_order_dict,
         )
         self.processed_text = ""
         self.citations: list[CitationInfo] = []

backend/danswer/chat/stream_processing/citation_processing.py

@@ -4,8 +4,8 @@ from collections.abc import Generator
 from danswer.chat.models import CitationInfo
 from danswer.chat.models import DanswerAnswerPiece
 from danswer.chat.models import LlmDoc
+from danswer.chat.stream_processing.utils import DocumentIdOrderMapping
 from danswer.configs.chat_configs import STOP_STREAM_PAT
-from danswer.llm.answering.stream_processing.utils import DocumentIdOrderMapping
 from danswer.prompts.constants import TRIPLE_BACKTICK
 from danswer.utils.logger import setup_logger
 
@@ -22,12 +22,16 @@ class CitationProcessor:
         self,
         context_docs: list[LlmDoc],
         doc_id_to_rank_map: DocumentIdOrderMapping,
+        display_doc_order_dict: dict[str, int],
         stop_stream: str | None = STOP_STREAM_PAT,
     ):
         self.context_docs = context_docs
         self.doc_id_to_rank_map = doc_id_to_rank_map
         self.stop_stream = stop_stream
         self.order_mapping = doc_id_to_rank_map.order_mapping
+        self.display_doc_order_dict = (
+            display_doc_order_dict  # original order of docs to displayed to user
+        )
         self.llm_out = ""
         self.max_citation_num = len(context_docs)
         self.citation_order: list[int] = []
@@ -98,6 +102,18 @@ class CitationProcessor:
                         self.citation_order.index(real_citation_num) + 1
                     )
 
+                    # get the value that was displayed to user, should always
+                    # be in the display_doc_order_dict. But check anyways
+                    if context_llm_doc.document_id in self.display_doc_order_dict:
+                        displayed_citation_num = self.display_doc_order_dict[
+                            context_llm_doc.document_id
+                        ]
+                    else:
+                        displayed_citation_num = real_citation_num
+                        logger.warning(
+                            f"Doc {context_llm_doc.document_id} not in display_doc_order_dict. Used LLM citation number instead."
+                        )
+
                     # Skip consecutive citations of the same work
                     if target_citation_num in self.current_citations:
                         start, end = citation.span()
@@ -118,6 +134,7 @@ class CitationProcessor:
                                 doc_id = int(match.group(1))
                                 context_llm_doc = self.context_docs[doc_id - 1]
                                 yield CitationInfo(
+                                    # stay with the original for now (order of LLM cites)
                                     citation_num=target_citation_num,
                                     document_id=context_llm_doc.document_id,
                                 )
@@ -139,6 +156,7 @@ class CitationProcessor:
                     if target_citation_num not in self.cited_inds:
                         self.cited_inds.add(target_citation_num)
                         yield CitationInfo(
+                            # stay with the original for now (order of LLM cites)
                             citation_num=target_citation_num,
                             document_id=context_llm_doc.document_id,
                         )
@@ -148,7 +166,8 @@ class CitationProcessor:
                         prev_length = len(self.curr_segment)
                         self.curr_segment = (
                             self.curr_segment[: start + length_to_add]
-                            + f"[[{target_citation_num}]]({link})"
+                            + f"[[{displayed_citation_num}]]({link})"  # use the value that was displayed to user
+                            # + f"[[{target_citation_num}]]({link})"
                             + self.curr_segment[end + length_to_add :]
                         )
                         length_to_add += len(self.curr_segment) - prev_length
@@ -156,7 +175,8 @@ class CitationProcessor:
                         prev_length = len(self.curr_segment)
                         self.curr_segment = (
                             self.curr_segment[: start + length_to_add]
-                            + f"[[{target_citation_num}]]()"
+                            + f"[[{displayed_citation_num}]]()"  # use the value that was displayed to user
+                            # + f"[[{target_citation_num}]]()"
                             + self.curr_segment[end + length_to_add :]
                         )
                         length_to_add += len(self.curr_segment) - prev_length

backend/danswer/chat/tool_handling/tool_response_handler.py

@@ -4,8 +4,8 @@ from langchain_core.messages import AIMessageChunk
 from langchain_core.messages import BaseMessage
 from langchain_core.messages import ToolCall
 
-from danswer.llm.answering.llm_response_handler import LLMCall
-from danswer.llm.answering.llm_response_handler import ResponsePart
+from danswer.chat.models import ResponsePart
+from danswer.chat.prompt_builder.build import LLMCall
 from danswer.llm.interfaces import LLM
 from danswer.tools.force import ForceUseTool
 from danswer.tools.message import build_tool_message

backend/danswer/configs/app_configs.py

@@ -43,9 +43,6 @@ WEB_DOMAIN = os.environ.get("WEB_DOMAIN") or "http://localhost:3000"
 AUTH_TYPE = AuthType((os.environ.get("AUTH_TYPE") or AuthType.DISABLED.value).lower())
 DISABLE_AUTH = AUTH_TYPE == AuthType.DISABLED
 
-# Necessary for cloud integration tests
-DISABLE_VERIFICATION = os.environ.get("DISABLE_VERIFICATION", "").lower() == "true"
-
 # Encryption key secret is used to encrypt connector credentials, api keys, and other sensitive
 # information. This provides an extra layer of security on top of Postgres access controls
 # and is available in Danswer EE
@@ -84,7 +81,14 @@ OAUTH_CLIENT_SECRET = (
     or ""
 )
 
+# for future OAuth connector support
+# OAUTH_CONFLUENCE_CLIENT_ID = os.environ.get("OAUTH_CONFLUENCE_CLIENT_ID", "")
+# OAUTH_CONFLUENCE_CLIENT_SECRET = os.environ.get("OAUTH_CONFLUENCE_CLIENT_SECRET", "")
+# OAUTH_JIRA_CLIENT_ID = os.environ.get("OAUTH_JIRA_CLIENT_ID", "")
+# OAUTH_JIRA_CLIENT_SECRET = os.environ.get("OAUTH_JIRA_CLIENT_SECRET", "")
+
 USER_AUTH_SECRET = os.environ.get("USER_AUTH_SECRET", "")
+
 # for basic auth
 REQUIRE_EMAIL_VERIFICATION = (
     os.environ.get("REQUIRE_EMAIL_VERIFICATION", "").lower() == "true"
@@ -118,6 +122,8 @@ VESPA_HOST = os.environ.get("VESPA_HOST") or "localhost"
 VESPA_CONFIG_SERVER_HOST = os.environ.get("VESPA_CONFIG_SERVER_HOST") or VESPA_HOST
 VESPA_PORT = os.environ.get("VESPA_PORT") or "8081"
 VESPA_TENANT_PORT = os.environ.get("VESPA_TENANT_PORT") or "19071"
+# the number of times to try and connect to vespa on startup before giving up
+VESPA_NUM_ATTEMPTS_ON_STARTUP = int(os.environ.get("NUM_RETRIES_ON_STARTUP") or 10)
 
 VESPA_CLOUD_URL = os.environ.get("VESPA_CLOUD_URL", "")
 
@@ -342,6 +348,12 @@ GITLAB_CONNECTOR_INCLUDE_CODE_FILES = (
     os.environ.get("GITLAB_CONNECTOR_INCLUDE_CODE_FILES", "").lower() == "true"
 )
 
+# Egnyte specific configs
+EGNYTE_LOCALHOST_OVERRIDE = os.getenv("EGNYTE_LOCALHOST_OVERRIDE")
+EGNYTE_BASE_DOMAIN = os.getenv("EGNYTE_DOMAIN")
+EGNYTE_CLIENT_ID = os.getenv("EGNYTE_CLIENT_ID")
+EGNYTE_CLIENT_SECRET = os.getenv("EGNYTE_CLIENT_SECRET")
+
 DASK_JOB_CLIENT_ENABLED = (
     os.environ.get("DASK_JOB_CLIENT_ENABLED", "").lower() == "true"
 )
@@ -405,21 +417,28 @@ LARGE_CHUNK_RATIO = 4
 # We don't want the metadata to overwhelm the actual contents of the chunk
 SKIP_METADATA_IN_CHUNK = os.environ.get("SKIP_METADATA_IN_CHUNK", "").lower() == "true"
 # Timeout to wait for job's last update before killing it, in hours
-CLEANUP_INDEXING_JOBS_TIMEOUT = int(os.environ.get("CLEANUP_INDEXING_JOBS_TIMEOUT", 3))
+CLEANUP_INDEXING_JOBS_TIMEOUT = int(
+    os.environ.get("CLEANUP_INDEXING_JOBS_TIMEOUT") or 3
+)
 
 # The indexer will warn in the logs whenver a document exceeds this threshold (in bytes)
 INDEXING_SIZE_WARNING_THRESHOLD = int(
-    os.environ.get("INDEXING_SIZE_WARNING_THRESHOLD", 100 * 1024 * 1024)
+    os.environ.get("INDEXING_SIZE_WARNING_THRESHOLD") or 100 * 1024 * 1024
 )
 
 # during indexing, will log verbose memory diff stats every x batches and at the end.
 # 0 disables this behavior and is the default.
-INDEXING_TRACER_INTERVAL = int(os.environ.get("INDEXING_TRACER_INTERVAL", 0))
+INDEXING_TRACER_INTERVAL = int(os.environ.get("INDEXING_TRACER_INTERVAL") or 0)
 
 # During an indexing attempt, specifies the number of batches which are allowed to
 # exception without aborting the attempt.
-INDEXING_EXCEPTION_LIMIT = int(os.environ.get("INDEXING_EXCEPTION_LIMIT", 0))
+INDEXING_EXCEPTION_LIMIT = int(os.environ.get("INDEXING_EXCEPTION_LIMIT") or 0)
 
+# Maximum file size in a document to be indexed
+MAX_DOCUMENT_CHARS = int(os.environ.get("MAX_DOCUMENT_CHARS") or 5_000_000)
+MAX_FILE_SIZE_BYTES = int(
+    os.environ.get("MAX_FILE_SIZE_BYTES") or 2 * 1024 * 1024 * 1024
+)  # 2GB in bytes
 
 #####
 # Miscellaneous

backend/danswer/configs/chat_configs.py

@@ -3,7 +3,6 @@ import os
 
 PROMPTS_YAML = "./danswer/seeding/prompts.yaml"
 PERSONAS_YAML = "./danswer/seeding/personas.yaml"
-INPUT_PROMPT_YAML = "./danswer/seeding/input_prompts.yaml"
 
 NUM_RETURNED_HITS = 50
 # Used for LLM filtering and reranking

backend/danswer/configs/constants.py

@@ -132,6 +132,7 @@ class DocumentSource(str, Enum):
     NOT_APPLICABLE = "not_applicable"
     FRESHDESK = "freshdesk"
     FIREFLIES = "fireflies"
+    EGNYTE = "egnyte"
 
 
 DocumentSourceRequiringTenantContext: list[DocumentSource] = [DocumentSource.FILE]

backend/danswer/configs/tool_configs.py

@@ -2,6 +2,8 @@ import json
 import os
 
 
+IMAGE_GENERATION_OUTPUT_FORMAT = os.environ.get("IMAGE_GENERATION_OUTPUT_FORMAT", "url")
+
 # if specified, will pass through request headers to the call to API calls made by custom tools
 CUSTOM_TOOL_PASS_THROUGH_HEADERS: list[str] | None = None
 _CUSTOM_TOOL_PASS_THROUGH_HEADERS_RAW = os.environ.get(

backend/danswer/connectors/confluence/connector.py

@@ -15,6 +15,7 @@ from danswer.connectors.confluence.utils import attachment_to_content
 from danswer.connectors.confluence.utils import build_confluence_document_id
 from danswer.connectors.confluence.utils import datetime_from_string
 from danswer.connectors.confluence.utils import extract_text_from_confluence_html
+from danswer.connectors.confluence.utils import validate_attachment_filetype
 from danswer.connectors.interfaces import GenerateDocumentsOutput
 from danswer.connectors.interfaces import GenerateSlimDocumentOutput
 from danswer.connectors.interfaces import LoadConnector
@@ -276,9 +277,11 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
         ):
             # If the page has restrictions, add them to the perm_sync_data
             # These will be used by doc_sync.py to sync permissions
-            perm_sync_data = {
-                "restrictions": page.get("restrictions", {}),
-                "space_key": page.get("space", {}).get("key"),
+            page_restrictions = page.get("restrictions")
+            page_space_key = page.get("space", {}).get("key")
+            page_perm_sync_data = {
+                "restrictions": page_restrictions or {},
+                "space_key": page_space_key,
             }
 
             doc_metadata_list.append(
@@ -288,7 +291,7 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
                         page["_links"]["webui"],
                         self.is_cloud,
                     ),
-                    perm_sync_data=perm_sync_data,
+                    perm_sync_data=page_perm_sync_data,
                 )
             )
             attachment_cql = f"type=attachment and container='{page['id']}'"
@@ -298,6 +301,21 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
                 expand=restrictions_expand,
                 limit=_SLIM_DOC_BATCH_SIZE,
             ):
+                if not validate_attachment_filetype(attachment):
+                    continue
+                attachment_restrictions = attachment.get("restrictions")
+                if not attachment_restrictions:
+                    attachment_restrictions = page_restrictions
+
+                attachment_space_key = attachment.get("space", {}).get("key")
+                if not attachment_space_key:
+                    attachment_space_key = page_space_key
+
+                attachment_perm_sync_data = {
+                    "restrictions": attachment_restrictions or {},
+                    "space_key": attachment_space_key,
+                }
+
                 doc_metadata_list.append(
                     SlimDocument(
                         id=build_confluence_document_id(
@@ -305,7 +323,7 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
                             attachment["_links"]["webui"],
                             self.is_cloud,
                         ),
-                        perm_sync_data=perm_sync_data,
+                        perm_sync_data=attachment_perm_sync_data,
                     )
                 )
             if len(doc_metadata_list) > _SLIM_DOC_BATCH_SIZE:

backend/danswer/connectors/confluence/onyx_confluence.py

@@ -368,4 +368,5 @@ def build_confluence_client(
         backoff_and_retry=True,
         max_backoff_retries=10,
         max_backoff_seconds=60,
+        cloud=is_cloud,
     )

backend/danswer/connectors/confluence/utils.py

@@ -177,19 +177,23 @@ def extract_text_from_confluence_html(
     return format_document_soup(soup)
 
 
-def attachment_to_content(
-    confluence_client: OnyxConfluence,
-    attachment: dict[str, Any],
-) -> str | None:
-    """If it returns None, assume that we should skip this attachment."""
-    if attachment["metadata"]["mediaType"] in [
+def validate_attachment_filetype(attachment: dict[str, Any]) -> bool:
+    return attachment["metadata"]["mediaType"] not in [
         "image/jpeg",
         "image/png",
         "image/gif",
         "image/svg+xml",
         "video/mp4",
         "video/quicktime",
-    ]:
+    ]
+
+
+def attachment_to_content(
+    confluence_client: OnyxConfluence,
+    attachment: dict[str, Any],
+) -> str | None:
+    """If it returns None, assume that we should skip this attachment."""
+    if not validate_attachment_filetype(attachment):
         return None
 
     download_link = confluence_client.url + attachment["_links"]["download"]
@@ -245,7 +249,7 @@ def build_confluence_document_id(
     return f"{base_url}{content_url}"
 
 
-def extract_referenced_attachment_names(page_text: str) -> list[str]:
+def _extract_referenced_attachment_names(page_text: str) -> list[str]:
     """Parse a Confluence html page to generate a list of current
         attachments in use
 

backend/danswer/connectors/egnyte/connector.py

@@ -0,0 +1,373 @@
+import io
+import os
+from collections.abc import Generator
+from datetime import datetime
+from datetime import timezone
+from logging import Logger
+from typing import Any
+from typing import cast
+from typing import IO
+
+import requests
+from retry import retry
+
+from danswer.configs.app_configs import EGNYTE_BASE_DOMAIN
+from danswer.configs.app_configs import EGNYTE_CLIENT_ID
+from danswer.configs.app_configs import EGNYTE_CLIENT_SECRET
+from danswer.configs.app_configs import EGNYTE_LOCALHOST_OVERRIDE
+from danswer.configs.app_configs import INDEX_BATCH_SIZE
+from danswer.configs.constants import DocumentSource
+from danswer.connectors.interfaces import GenerateDocumentsOutput
+from danswer.connectors.interfaces import LoadConnector
+from danswer.connectors.interfaces import OAuthConnector
+from danswer.connectors.interfaces import PollConnector
+from danswer.connectors.interfaces import SecondsSinceUnixEpoch
+from danswer.connectors.models import BasicExpertInfo
+from danswer.connectors.models import ConnectorMissingCredentialError
+from danswer.connectors.models import Document
+from danswer.connectors.models import Section
+from danswer.file_processing.extract_file_text import detect_encoding
+from danswer.file_processing.extract_file_text import extract_file_text
+from danswer.file_processing.extract_file_text import get_file_ext
+from danswer.file_processing.extract_file_text import is_text_file_extension
+from danswer.file_processing.extract_file_text import is_valid_file_ext
+from danswer.file_processing.extract_file_text import read_text_file
+from danswer.utils.logger import setup_logger
+
+
+logger = setup_logger()
+
+_EGNYTE_API_BASE = "https://{domain}.egnyte.com/pubapi/v1"
+_EGNYTE_APP_BASE = "https://{domain}.egnyte.com"
+_TIMEOUT = 60
+
+
+def _request_with_retries(
+    method: str,
+    url: str,
+    data: dict[str, Any] | None = None,
+    headers: dict[str, Any] | None = None,
+    params: dict[str, Any] | None = None,
+    timeout: int = _TIMEOUT,
+    stream: bool = False,
+    tries: int = 8,
+    delay: float = 1,
+    backoff: float = 2,
+) -> requests.Response:
+    @retry(tries=tries, delay=delay, backoff=backoff, logger=cast(Logger, logger))
+    def _make_request() -> requests.Response:
+        response = requests.request(
+            method,
+            url,
+            data=data,
+            headers=headers,
+            params=params,
+            timeout=timeout,
+            stream=stream,
+        )
+        response.raise_for_status()
+        return response
+
+    return _make_request()
+
+
+def _parse_last_modified(last_modified: str) -> datetime:
+    return datetime.strptime(last_modified, "%a, %d %b %Y %H:%M:%S %Z").replace(
+        tzinfo=timezone.utc
+    )
+
+
+def _process_egnyte_file(
+    file_metadata: dict[str, Any],
+    file_content: IO,
+    base_url: str,
+    folder_path: str | None = None,
+) -> Document | None:
+    """Process an Egnyte file into a Document object
+
+    Args:
+        file_data: The file data from Egnyte API
+        file_content: The raw content of the file in bytes
+        base_url: The base URL for the Egnyte instance
+        folder_path: Optional folder path to filter results
+    """
+    # Skip if file path doesn't match folder path filter
+    if folder_path and not file_metadata["path"].startswith(folder_path):
+        raise ValueError(
+            f"File path {file_metadata['path']} does not match folder path {folder_path}"
+        )
+
+    file_name = file_metadata["name"]
+    extension = get_file_ext(file_name)
+    if not is_valid_file_ext(extension):
+        logger.warning(f"Skipping file '{file_name}' with extension '{extension}'")
+        return None
+
+    # Extract text content based on file type
+    if is_text_file_extension(file_name):
+        encoding = detect_encoding(file_content)
+        file_content_raw, file_metadata = read_text_file(
+            file_content, encoding=encoding, ignore_danswer_metadata=False
+        )
+    else:
+        file_content_raw = extract_file_text(
+            file=file_content,
+            file_name=file_name,
+            break_on_unprocessable=True,
+        )
+
+    # Build the web URL for the file
+    web_url = f"{base_url}/navigate/file/{file_metadata['group_id']}"
+
+    # Create document metadata
+    metadata: dict[str, str | list[str]] = {
+        "file_path": file_metadata["path"],
+        "last_modified": file_metadata.get("last_modified", ""),
+    }
+
+    # Add lock info if present
+    if lock_info := file_metadata.get("lock_info"):
+        metadata[
+            "lock_owner"
+        ] = f"{lock_info.get('first_name', '')} {lock_info.get('last_name', '')}"
+
+    # Create the document owners
+    primary_owner = None
+    if uploaded_by := file_metadata.get("uploaded_by"):
+        primary_owner = BasicExpertInfo(
+            email=uploaded_by,  # Using username as email since that's what we have
+        )
+
+    # Create the document
+    return Document(
+        id=f"egnyte-{file_metadata['entry_id']}",
+        sections=[Section(text=file_content_raw.strip(), link=web_url)],
+        source=DocumentSource.EGNYTE,
+        semantic_identifier=file_name,
+        metadata=metadata,
+        doc_updated_at=(
+            _parse_last_modified(file_metadata["last_modified"])
+            if "last_modified" in file_metadata
+            else None
+        ),
+        primary_owners=[primary_owner] if primary_owner else None,
+    )
+
+
+class EgnyteConnector(LoadConnector, PollConnector, OAuthConnector):
+    def __init__(
+        self,
+        folder_path: str | None = None,
+        batch_size: int = INDEX_BATCH_SIZE,
+    ) -> None:
+        self.domain = ""  # will always be set in `load_credentials`
+        self.folder_path = folder_path or ""  # Root folder if not specified
+        self.batch_size = batch_size
+        self.access_token: str | None = None
+
+    @classmethod
+    def oauth_id(cls) -> DocumentSource:
+        return DocumentSource.EGNYTE
+
+    @classmethod
+    def oauth_authorization_url(cls, base_domain: str, state: str) -> str:
+        if not EGNYTE_CLIENT_ID:
+            raise ValueError("EGNYTE_CLIENT_ID environment variable must be set")
+        if not EGNYTE_BASE_DOMAIN:
+            raise ValueError("EGNYTE_DOMAIN environment variable must be set")
+
+        if EGNYTE_LOCALHOST_OVERRIDE:
+            base_domain = EGNYTE_LOCALHOST_OVERRIDE
+
+        callback_uri = f"{base_domain.strip('/')}/connector/oauth/callback/egnyte"
+        return (
+            f"https://{EGNYTE_BASE_DOMAIN}.egnyte.com/puboauth/token"
+            f"?client_id={EGNYTE_CLIENT_ID}"
+            f"&redirect_uri={callback_uri}"
+            f"&scope=Egnyte.filesystem"
+            f"&state={state}"
+            f"&response_type=code"
+        )
+
+    @classmethod
+    def oauth_code_to_token(cls, code: str) -> dict[str, Any]:
+        if not EGNYTE_CLIENT_ID:
+            raise ValueError("EGNYTE_CLIENT_ID environment variable must be set")
+        if not EGNYTE_CLIENT_SECRET:
+            raise ValueError("EGNYTE_CLIENT_SECRET environment variable must be set")
+        if not EGNYTE_BASE_DOMAIN:
+            raise ValueError("EGNYTE_DOMAIN environment variable must be set")
+
+        # Exchange code for token
+        url = f"https://{EGNYTE_BASE_DOMAIN}.egnyte.com/puboauth/token"
+        data = {
+            "client_id": EGNYTE_CLIENT_ID,
+            "client_secret": EGNYTE_CLIENT_SECRET,
+            "code": code,
+            "grant_type": "authorization_code",
+            "redirect_uri": f"{EGNYTE_LOCALHOST_OVERRIDE or ''}/connector/oauth/callback/egnyte",
+            "scope": "Egnyte.filesystem",
+        }
+        headers = {"Content-Type": "application/x-www-form-urlencoded"}
+
+        response = _request_with_retries(
+            method="POST",
+            url=url,
+            data=data,
+            headers=headers,
+            # try a lot faster since this is a realtime flow
+            backoff=0,
+            delay=0.1,
+        )
+        if not response.ok:
+            raise RuntimeError(f"Failed to exchange code for token: {response.text}")
+
+        token_data = response.json()
+        return {
+            "domain": EGNYTE_BASE_DOMAIN,
+            "access_token": token_data["access_token"],
+        }
+
+    def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
+        self.domain = credentials["domain"]
+        self.access_token = credentials["access_token"]
+        return None
+
+    def _get_files_list(
+        self,
+        path: str,
+    ) -> list[dict[str, Any]]:
+        if not self.access_token or not self.domain:
+            raise ConnectorMissingCredentialError("Egnyte")
+
+        headers = {
+            "Authorization": f"Bearer {self.access_token}",
+        }
+
+        params: dict[str, Any] = {
+            "list_content": True,
+        }
+
+        url = f"{_EGNYTE_API_BASE.format(domain=self.domain)}/fs/{path or ''}"
+        response = _request_with_retries(
+            method="GET", url=url, headers=headers, params=params, timeout=_TIMEOUT
+        )
+        if not response.ok:
+            raise RuntimeError(f"Failed to fetch files from Egnyte: {response.text}")
+
+        data = response.json()
+        all_files: list[dict[str, Any]] = []
+
+        # Add files from current directory
+        all_files.extend(data.get("files", []))
+
+        # Recursively traverse folders
+        for item in data.get("folders", []):
+            all_files.extend(self._get_files_list(item["path"]))
+
+        return all_files
+
+    def _filter_files(
+        self,
+        files: list[dict[str, Any]],
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+    ) -> list[dict[str, Any]]:
+        filtered_files = []
+        for file in files:
+            if file["is_folder"]:
+                continue
+
+            file_modified = _parse_last_modified(file["last_modified"])
+            if start_time and file_modified < start_time:
+                continue
+            if end_time and file_modified > end_time:
+                continue
+
+            filtered_files.append(file)
+
+        return filtered_files
+
+    def _process_files(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+    ) -> Generator[list[Document], None, None]:
+        files = self._get_files_list(self.folder_path)
+        files = self._filter_files(files, start_time, end_time)
+
+        current_batch: list[Document] = []
+        for file in files:
+            try:
+                # Set up request with streaming enabled
+                headers = {
+                    "Authorization": f"Bearer {self.access_token}",
+                }
+                url = f"{_EGNYTE_API_BASE.format(domain=self.domain)}/fs-content/{file['path']}"
+                response = _request_with_retries(
+                    method="GET",
+                    url=url,
+                    headers=headers,
+                    timeout=_TIMEOUT,
+                    stream=True,
+                )
+
+                if not response.ok:
+                    logger.error(
+                        f"Failed to fetch file content: {file['path']} (status code: {response.status_code})"
+                    )
+                    continue
+
+                # Stream the response content into a BytesIO buffer
+                buffer = io.BytesIO()
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        buffer.write(chunk)
+
+                # Reset buffer's position to the start
+                buffer.seek(0)
+
+                # Process the streamed file content
+                doc = _process_egnyte_file(
+                    file_metadata=file,
+                    file_content=buffer,
+                    base_url=_EGNYTE_APP_BASE.format(domain=self.domain),
+                    folder_path=self.folder_path,
+                )
+
+                if doc is not None:
+                    current_batch.append(doc)
+
+                    if len(current_batch) >= self.batch_size:
+                        yield current_batch
+                        current_batch = []
+
+            except Exception:
+                logger.exception(f"Failed to process file {file['path']}")
+                continue
+
+        if current_batch:
+            yield current_batch
+
+    def load_from_state(self) -> GenerateDocumentsOutput:
+        yield from self._process_files()
+
+    def poll_source(
+        self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch
+    ) -> GenerateDocumentsOutput:
+        start_time = datetime.fromtimestamp(start, tz=timezone.utc)
+        end_time = datetime.fromtimestamp(end, tz=timezone.utc)
+
+        yield from self._process_files(start_time=start_time, end_time=end_time)
+
+
+if __name__ == "__main__":
+    connector = EgnyteConnector()
+    connector.load_credentials(
+        {
+            "domain": os.environ["EGNYTE_DOMAIN"],
+            "access_token": os.environ["EGNYTE_ACCESS_TOKEN"],
+        }
+    )
+    document_batches = connector.load_from_state()
+    print(next(document_batches))

backend/danswer/connectors/factory.py

@@ -15,6 +15,7 @@ from danswer.connectors.danswer_jira.connector import JiraConnector
 from danswer.connectors.discourse.connector import DiscourseConnector
 from danswer.connectors.document360.connector import Document360Connector
 from danswer.connectors.dropbox.connector import DropboxConnector
+from danswer.connectors.egnyte.connector import EgnyteConnector
 from danswer.connectors.file.connector import LocalFileConnector
 from danswer.connectors.fireflies.connector import FirefliesConnector
 from danswer.connectors.freshdesk.connector import FreshdeskConnector
@@ -103,6 +104,7 @@ def identify_connector_class(
         DocumentSource.XENFORO: XenforoConnector,
         DocumentSource.FRESHDESK: FreshdeskConnector,
         DocumentSource.FIREFLIES: FirefliesConnector,
+        DocumentSource.EGNYTE: EgnyteConnector,
     }
     connector_by_source = connector_map.get(source, {})
 

backend/danswer/connectors/file/connector.py

@@ -17,11 +17,11 @@ from danswer.connectors.models import BasicExpertInfo
 from danswer.connectors.models import Document
 from danswer.connectors.models import Section
 from danswer.db.engine import get_session_with_tenant
-from danswer.file_processing.extract_file_text import check_file_ext_is_valid
 from danswer.file_processing.extract_file_text import detect_encoding
 from danswer.file_processing.extract_file_text import extract_file_text
 from danswer.file_processing.extract_file_text import get_file_ext
 from danswer.file_processing.extract_file_text import is_text_file_extension
+from danswer.file_processing.extract_file_text import is_valid_file_ext
 from danswer.file_processing.extract_file_text import load_files_from_zip
 from danswer.file_processing.extract_file_text import read_pdf_file
 from danswer.file_processing.extract_file_text import read_text_file
@@ -50,7 +50,7 @@ def _read_files_and_metadata(
             file_content, ignore_dirs=True
         ):
             yield os.path.join(directory_path, file_info.filename), file, metadata
-    elif check_file_ext_is_valid(extension):
+    elif is_valid_file_ext(extension):
         yield file_name, file_content, metadata
     else:
         logger.warning(f"Skipping file '{file_name}' with extension '{extension}'")
@@ -63,7 +63,7 @@ def _process_file(
     pdf_pass: str | None = None,
 ) -> list[Document]:
     extension = get_file_ext(file_name)
-    if not check_file_ext_is_valid(extension):
+    if not is_valid_file_ext(extension):
         logger.warning(f"Skipping file '{file_name}' with extension '{extension}'")
         return []
 

backend/danswer/connectors/google_drive/connector.py

@@ -4,11 +4,13 @@ from concurrent.futures import as_completed
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from typing import Any
+from typing import cast
 
 from google.oauth2.credentials import Credentials as OAuthCredentials  # type: ignore
 from google.oauth2.service_account import Credentials as ServiceAccountCredentials  # type: ignore
 
 from danswer.configs.app_configs import INDEX_BATCH_SIZE
+from danswer.configs.app_configs import MAX_FILE_SIZE_BYTES
 from danswer.configs.constants import DocumentSource
 from danswer.connectors.google_drive.doc_conversion import build_slim_document
 from danswer.connectors.google_drive.doc_conversion import (
@@ -452,12 +454,14 @@ class GoogleDriveConnector(LoadConnector, PollConnector, SlimConnector):
             if isinstance(self.creds, ServiceAccountCredentials)
             else self._manage_oauth_retrieval
         )
-        return retrieval_method(
+        drive_files = retrieval_method(
             is_slim=is_slim,
             start=start,
             end=end,
         )
 
+        return drive_files
+
     def _extract_docs_from_google_drive(
         self,
         start: SecondsSinceUnixEpoch | None = None,
@@ -473,6 +477,15 @@ class GoogleDriveConnector(LoadConnector, PollConnector, SlimConnector):
         files_to_process = []
         # Gather the files into batches to be processed in parallel
         for file in self._fetch_drive_items(is_slim=False, start=start, end=end):
+            if (
+                file.get("size")
+                and int(cast(str, file.get("size"))) > MAX_FILE_SIZE_BYTES
+            ):
+                logger.warning(
+                    f"Skipping file {file.get('name', 'Unknown')} as it is too large: {file.get('size')} bytes"
+                )
+                continue
+
             files_to_process.append(file)
             if len(files_to_process) >= LARGE_BATCH_SIZE:
                 yield from _process_files_batch(

backend/danswer/connectors/google_drive/file_retrieval.py

@@ -16,7 +16,7 @@ logger = setup_logger()
 
 FILE_FIELDS = (
     "nextPageToken, files(mimeType, id, name, permissions, modifiedTime, webViewLink, "
-    "shortcutDetails, owners(emailAddress))"
+    "shortcutDetails, owners(emailAddress), size)"
 )
 SLIM_FILE_FIELDS = (
     "nextPageToken, files(mimeType, id, name, permissions(emailAddress, type), "

backend/danswer/connectors/interfaces.py

@@ -2,6 +2,7 @@ import abc
 from collections.abc import Iterator
 from typing import Any
 
+from danswer.configs.constants import DocumentSource
 from danswer.connectors.models import Document
 from danswer.connectors.models import SlimDocument
 
@@ -64,6 +65,23 @@ class SlimConnector(BaseConnector):
         raise NotImplementedError
 
 
+class OAuthConnector(BaseConnector):
+    @classmethod
+    @abc.abstractmethod
+    def oauth_id(cls) -> DocumentSource:
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def oauth_authorization_url(cls, base_domain: str, state: str) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def oauth_code_to_token(cls, code: str) -> dict[str, Any]:
+        raise NotImplementedError
+
+
 # Event driven
 class EventConnector(BaseConnector):
     @abc.abstractmethod

backend/danswer/connectors/linear/connector.py

@@ -132,7 +132,6 @@ class LinearConnector(LoadConnector, PollConnector):
                             branchName
                             customerTicketCount
                             description
-                            descriptionData
                             comments {
                                 nodes {
                                     url
@@ -215,5 +214,6 @@ class LinearConnector(LoadConnector, PollConnector):
 if __name__ == "__main__":
     connector = LinearConnector()
     connector.load_credentials({"linear_api_key": os.environ["LINEAR_API_KEY"]})
+
     document_batches = connector.load_from_state()
     print(next(document_batches))

backend/danswer/connectors/slack/connector.py

@@ -171,7 +171,9 @@ def thread_to_doc(
         else first_message
     )
 
-    doc_sem_id = f"{initial_sender_name} in #{channel['name']}: {snippet}"
+    doc_sem_id = f"{initial_sender_name} in #{channel['name']}: {snippet}".replace(
+        "\n", " "
+    )
 
     return Document(
         id=f"{channel_id}__{thread[0]['ts']}",

backend/danswer/connectors/teams/connector.py

@@ -33,7 +33,7 @@ def get_created_datetime(chat_message: ChatMessage) -> datetime:
 
 def _extract_channel_members(channel: Channel) -> list[BasicExpertInfo]:
     channel_members_list: list[BasicExpertInfo] = []
-    members = channel.members.get().execute_query()
+    members = channel.members.get().execute_query_retry()
     for member in members:
         channel_members_list.append(BasicExpertInfo(display_name=member.display_name))
     return channel_members_list
@@ -51,7 +51,7 @@ def _get_threads_from_channel(
         end = end.replace(tzinfo=timezone.utc)
 
     query = channel.messages.get()
-    base_messages: list[ChatMessage] = query.execute_query()
+    base_messages: list[ChatMessage] = query.execute_query_retry()
 
     threads: list[list[ChatMessage]] = []
     for base_message in base_messages:
@@ -65,7 +65,7 @@ def _get_threads_from_channel(
             continue
 
         reply_query = base_message.replies.get_all()
-        replies = reply_query.execute_query()
+        replies = reply_query.execute_query_retry()
 
         # start a list containing the base message and its replies
         thread: list[ChatMessage] = [base_message]
@@ -82,7 +82,7 @@ def _get_channels_from_teams(
     channels_list: list[Channel] = []
     for team in teams:
         query = team.channels.get()
-        channels = query.execute_query()
+        channels = query.execute_query_retry()
         channels_list.extend(channels)
 
     return channels_list
@@ -210,7 +210,7 @@ class TeamsConnector(LoadConnector, PollConnector):
 
         teams_list: list[Team] = []
 
-        teams = self.graph_client.teams.get().execute_query()
+        teams = self.graph_client.teams.get().execute_query_retry()
 
         if len(self.requested_team_list) > 0:
             adjusted_request_strings = [
@@ -234,14 +234,25 @@ class TeamsConnector(LoadConnector, PollConnector):
             raise ConnectorMissingCredentialError("Teams")
 
         teams = self._get_all_teams()
+        logger.debug(f"Found available teams: {[str(t) for t in teams]}")
+        if not teams:
+            msg = "No teams found."
+            logger.error(msg)
+            raise ValueError(msg)
 
         channels = _get_channels_from_teams(
             teams=teams,
         )
+        logger.debug(f"Found available channels: {[c.id for c in channels]}")
+        if not channels:
+            msg = "No channels found."
+            logger.error(msg)
+            raise ValueError(msg)
 
         # goes over channels, converts them into Document objects and then yields them in batches
         doc_batch: list[Document] = []
         for channel in channels:
+            logger.debug(f"Fetching threads from channel: {channel.id}")
             thread_list = _get_threads_from_channel(channel, start=start, end=end)
             for thread in thread_list:
                 converted_doc = _convert_thread_to_document(channel, thread)
@@ -259,8 +270,8 @@ class TeamsConnector(LoadConnector, PollConnector):
     def poll_source(
         self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch
     ) -> GenerateDocumentsOutput:
-        start_datetime = datetime.utcfromtimestamp(start)
-        end_datetime = datetime.utcfromtimestamp(end)
+        start_datetime = datetime.fromtimestamp(start, timezone.utc)
+        end_datetime = datetime.fromtimestamp(end, timezone.utc)
         return self._fetch_from_teams(start=start_datetime, end=end_datetime)
 
 

backend/danswer/context/search/pipeline.py

@@ -5,7 +5,11 @@ from typing import cast
 
 from sqlalchemy.orm import Session
 
+from danswer.chat.models import PromptConfig
 from danswer.chat.models import SectionRelevancePiece
+from danswer.chat.prune_and_merge import _merge_sections
+from danswer.chat.prune_and_merge import ChunkRange
+from danswer.chat.prune_and_merge import merge_chunk_intervals
 from danswer.configs.chat_configs import DISABLE_LLM_DOC_RELEVANCE
 from danswer.context.search.enums import LLMEvaluationType
 from danswer.context.search.enums import QueryFlow
@@ -27,10 +31,6 @@ from danswer.db.models import User
 from danswer.db.search_settings import get_current_search_settings
 from danswer.document_index.factory import get_default_document_index
 from danswer.document_index.interfaces import VespaChunkRequest
-from danswer.llm.answering.models import PromptConfig
-from danswer.llm.answering.prune_and_merge import _merge_sections
-from danswer.llm.answering.prune_and_merge import ChunkRange
-from danswer.llm.answering.prune_and_merge import merge_chunk_intervals
 from danswer.llm.interfaces import LLM
 from danswer.secondary_llm_flows.agentic_evaluation import evaluate_inference_section
 from danswer.utils.logger import setup_logger

backend/danswer/danswerbot/slack/blocks.py

@@ -204,7 +204,8 @@ def _build_documents_blocks(
             continue
         seen_docs_identifiers.add(d.document_id)
 
-        doc_sem_id = d.semantic_identifier
+        # Strip newlines from the semantic identifier for Slackbot formatting
+        doc_sem_id = d.semantic_identifier.replace("\n", " ")
         if d.source_type == DocumentSource.SLACK.value:
             doc_sem_id = "#" + doc_sem_id
 

backend/danswer/danswerbot/slack/handlers/handle_regular_answer.py

@@ -373,7 +373,9 @@ def handle_regular_answer(
         respond_in_thread(
             client=client,
             channel=channel,
-            receiver_ids=receiver_ids,
+            receiver_ids=[message_info.sender]
+            if message_info.is_bot_msg and message_info.sender
+            else receiver_ids,
             text="Hello! Danswer has some results for you!",
             blocks=all_blocks,
             thread_ts=message_ts_to_respond_to,

backend/danswer/danswerbot/slack/utils.py

@@ -11,6 +11,7 @@ from retry import retry
 from slack_sdk import WebClient
 from slack_sdk.errors import SlackApiError
 from slack_sdk.models.blocks import Block
+from slack_sdk.models.blocks import SectionBlock
 from slack_sdk.models.metadata import Metadata
 from slack_sdk.socket_mode import SocketModeClient
 
@@ -140,6 +141,40 @@ def remove_danswer_bot_tag(message_str: str, client: WebClient) -> str:
     return re.sub(rf"<@{bot_tag_id}>\s", "", message_str)
 
 
+def _check_for_url_in_block(block: Block) -> bool:
+    """
+    Check if the block has a key that contains "url" in it
+    """
+    block_dict = block.to_dict()
+
+    def check_dict_for_url(d: dict) -> bool:
+        for key, value in d.items():
+            if "url" in key.lower():
+                return True
+            if isinstance(value, dict):
+                if check_dict_for_url(value):
+                    return True
+            elif isinstance(value, list):
+                for item in value:
+                    if isinstance(item, dict) and check_dict_for_url(item):
+                        return True
+        return False
+
+    return check_dict_for_url(block_dict)
+
+
+def _build_error_block(error_message: str) -> Block:
+    """
+    Build an error block to display in slack so that the user can see
+    the error without completely breaking
+    """
+    display_text = (
+        "There was an error displaying all of the Onyx answers."
+        f" Please let an admin or an onyx developer know. Error: {error_message}"
+    )
+    return SectionBlock(text=display_text)
+
+
 @retry(
     tries=DANSWER_BOT_NUM_RETRIES,
     delay=0.25,
@@ -162,24 +197,9 @@ def respond_in_thread(
     message_ids: list[str] = []
     if not receiver_ids:
         slack_call = make_slack_api_rate_limited(client.chat_postMessage)
-        response = slack_call(
-            channel=channel,
-            text=text,
-            blocks=blocks,
-            thread_ts=thread_ts,
-            metadata=metadata,
-            unfurl_links=unfurl,
-            unfurl_media=unfurl,
-        )
-        if not response.get("ok"):
-            raise RuntimeError(f"Failed to post message: {response}")
-        message_ids.append(response["message_ts"])
-    else:
-        slack_call = make_slack_api_rate_limited(client.chat_postEphemeral)
-        for receiver in receiver_ids:
+        try:
             response = slack_call(
                 channel=channel,
-                user=receiver,
                 text=text,
                 blocks=blocks,
                 thread_ts=thread_ts,
@@ -187,8 +207,68 @@ def respond_in_thread(
                 unfurl_links=unfurl,
                 unfurl_media=unfurl,
             )
-            if not response.get("ok"):
-                raise RuntimeError(f"Failed to post message: {response}")
+        except Exception as e:
+            logger.warning(f"Failed to post message: {e} \n blocks: {blocks}")
+            logger.warning("Trying again without blocks that have urls")
+
+            if not blocks:
+                raise e
+
+            blocks_without_urls = [
+                block for block in blocks if not _check_for_url_in_block(block)
+            ]
+            blocks_without_urls.append(_build_error_block(str(e)))
+
+            # Try again wtihout blocks containing url
+            response = slack_call(
+                channel=channel,
+                text=text,
+                blocks=blocks_without_urls,
+                thread_ts=thread_ts,
+                metadata=metadata,
+                unfurl_links=unfurl,
+                unfurl_media=unfurl,
+            )
+
+        message_ids.append(response["message_ts"])
+    else:
+        slack_call = make_slack_api_rate_limited(client.chat_postEphemeral)
+        for receiver in receiver_ids:
+            try:
+                response = slack_call(
+                    channel=channel,
+                    user=receiver,
+                    text=text,
+                    blocks=blocks,
+                    thread_ts=thread_ts,
+                    metadata=metadata,
+                    unfurl_links=unfurl,
+                    unfurl_media=unfurl,
+                )
+            except Exception as e:
+                logger.warning(f"Failed to post message: {e} \n blocks: {blocks}")
+                logger.warning("Trying again without blocks that have urls")
+
+                if not blocks:
+                    raise e
+
+                blocks_without_urls = [
+                    block for block in blocks if not _check_for_url_in_block(block)
+                ]
+                blocks_without_urls.append(_build_error_block(str(e)))
+
+                # Try again wtihout blocks containing url
+                response = slack_call(
+                    channel=channel,
+                    user=receiver,
+                    text=text,
+                    blocks=blocks_without_urls,
+                    thread_ts=thread_ts,
+                    metadata=metadata,
+                    unfurl_links=unfurl,
+                    unfurl_media=unfurl,
+                )
+
             message_ids.append(response["message_ts"])
 
     return message_ids

backend/danswer/db/credentials.py

@@ -20,7 +20,6 @@ from danswer.db.models import DocumentByConnectorCredentialPair
 from danswer.db.models import User
 from danswer.db.models import User__UserGroup
 from danswer.server.documents.models import CredentialBase
-from danswer.server.documents.models import CredentialDataUpdateRequest
 from danswer.utils.logger import setup_logger
 
 
@@ -248,7 +247,6 @@ def create_credential(
     )
 
     db_session.commit()
-
     return credential
 
 
@@ -263,7 +261,8 @@ def _cleanup_credential__user_group_relationships__no_commit(
 
 def alter_credential(
     credential_id: int,
-    credential_data: CredentialDataUpdateRequest,
+    name: str,
+    credential_json: dict[str, Any],
     user: User,
     db_session: Session,
 ) -> Credential | None:
@@ -273,11 +272,13 @@ def alter_credential(
     if credential is None:
         return None
 
-    credential.name = credential_data.name
+    credential.name = name
 
-    # Update only the keys present in credential_data.credential_json
-    for key, value in credential_data.credential_json.items():
-        credential.credential_json[key] = value
+    # Assign a new dictionary to credential.credential_json
+    credential.credential_json = {
+        **credential.credential_json,
+        **credential_json,
+    }
 
     credential.user_id = user.id if user is not None else None
     db_session.commit()
@@ -310,8 +311,8 @@ def update_credential_json(
     credential = fetch_credential_by_id(credential_id, user, db_session)
     if credential is None:
         return None
-    credential.credential_json = credential_json
 
+    credential.credential_json = credential_json
     db_session.commit()
     return credential
 

backend/danswer/db/index_attempt.py

@@ -522,12 +522,16 @@ def expire_index_attempts(
     search_settings_id: int,
     db_session: Session,
 ) -> None:
-    delete_query = (
-        delete(IndexAttempt)
+    not_started_query = (
+        update(IndexAttempt)
         .where(IndexAttempt.search_settings_id == search_settings_id)
         .where(IndexAttempt.status == IndexingStatus.NOT_STARTED)
+        .values(
+            status=IndexingStatus.CANCELED,
+            error_msg="Canceled, likely due to model swap",
+        )
     )
-    db_session.execute(delete_query)
+    db_session.execute(not_started_query)
 
     update_query = (
         update(IndexAttempt)
@@ -549,9 +553,14 @@ def cancel_indexing_attempts_for_ccpair(
     include_secondary_index: bool = False,
 ) -> None:
     stmt = (
-        delete(IndexAttempt)
+        update(IndexAttempt)
         .where(IndexAttempt.connector_credential_pair_id == cc_pair_id)
         .where(IndexAttempt.status == IndexingStatus.NOT_STARTED)
+        .values(
+            status=IndexingStatus.CANCELED,
+            error_msg="Canceled by user",
+            time_started=datetime.now(timezone.utc),
+        )
     )
 
     if not include_secondary_index:

backend/danswer/db/input_prompt.py

@@ -1,202 +0,0 @@
-from uuid import UUID
-
-from fastapi import HTTPException
-from sqlalchemy import select
-from sqlalchemy.orm import Session
-
-from danswer.db.models import InputPrompt
-from danswer.db.models import User
-from danswer.server.features.input_prompt.models import InputPromptSnapshot
-from danswer.server.manage.models import UserInfo
-from danswer.utils.logger import setup_logger
-
-
-logger = setup_logger()
-
-
-def insert_input_prompt_if_not_exists(
-    user: User | None,
-    input_prompt_id: int | None,
-    prompt: str,
-    content: str,
-    active: bool,
-    is_public: bool,
-    db_session: Session,
-    commit: bool = True,
-) -> InputPrompt:
-    if input_prompt_id is not None:
-        input_prompt = (
-            db_session.query(InputPrompt).filter_by(id=input_prompt_id).first()
-        )
-    else:
-        query = db_session.query(InputPrompt).filter(InputPrompt.prompt == prompt)
-        if user:
-            query = query.filter(InputPrompt.user_id == user.id)
-        else:
-            query = query.filter(InputPrompt.user_id.is_(None))
-        input_prompt = query.first()
-
-    if input_prompt is None:
-        input_prompt = InputPrompt(
-            id=input_prompt_id,
-            prompt=prompt,
-            content=content,
-            active=active,
-            is_public=is_public or user is None,
-            user_id=user.id if user else None,
-        )
-        db_session.add(input_prompt)
-
-    if commit:
-        db_session.commit()
-
-    return input_prompt
-
-
-def insert_input_prompt(
-    prompt: str,
-    content: str,
-    is_public: bool,
-    user: User | None,
-    db_session: Session,
-) -> InputPrompt:
-    input_prompt = InputPrompt(
-        prompt=prompt,
-        content=content,
-        active=True,
-        is_public=is_public or user is None,
-        user_id=user.id if user is not None else None,
-    )
-    db_session.add(input_prompt)
-    db_session.commit()
-
-    return input_prompt
-
-
-def update_input_prompt(
-    user: User | None,
-    input_prompt_id: int,
-    prompt: str,
-    content: str,
-    active: bool,
-    db_session: Session,
-) -> InputPrompt:
-    input_prompt = db_session.scalar(
-        select(InputPrompt).where(InputPrompt.id == input_prompt_id)
-    )
-    if input_prompt is None:
-        raise ValueError(f"No input prompt with id {input_prompt_id}")
-
-    if not validate_user_prompt_authorization(user, input_prompt):
-        raise HTTPException(status_code=401, detail="You don't own this prompt")
-
-    input_prompt.prompt = prompt
-    input_prompt.content = content
-    input_prompt.active = active
-
-    db_session.commit()
-    return input_prompt
-
-
-def validate_user_prompt_authorization(
-    user: User | None, input_prompt: InputPrompt
-) -> bool:
-    prompt = InputPromptSnapshot.from_model(input_prompt=input_prompt)
-
-    if prompt.user_id is not None:
-        if user is None:
-            return False
-
-        user_details = UserInfo.from_model(user)
-        if str(user_details.id) != str(prompt.user_id):
-            return False
-    return True
-
-
-def remove_public_input_prompt(input_prompt_id: int, db_session: Session) -> None:
-    input_prompt = db_session.scalar(
-        select(InputPrompt).where(InputPrompt.id == input_prompt_id)
-    )
-
-    if input_prompt is None:
-        raise ValueError(f"No input prompt with id {input_prompt_id}")
-
-    if not input_prompt.is_public:
-        raise HTTPException(status_code=400, detail="This prompt is not public")
-
-    db_session.delete(input_prompt)
-    db_session.commit()
-
-
-def remove_input_prompt(
-    user: User | None, input_prompt_id: int, db_session: Session
-) -> None:
-    input_prompt = db_session.scalar(
-        select(InputPrompt).where(InputPrompt.id == input_prompt_id)
-    )
-    if input_prompt is None:
-        raise ValueError(f"No input prompt with id {input_prompt_id}")
-
-    if input_prompt.is_public:
-        raise HTTPException(
-            status_code=400, detail="Cannot delete public prompts with this method"
-        )
-
-    if not validate_user_prompt_authorization(user, input_prompt):
-        raise HTTPException(status_code=401, detail="You do not own this prompt")
-
-    db_session.delete(input_prompt)
-    db_session.commit()
-
-
-def fetch_input_prompt_by_id(
-    id: int, user_id: UUID | None, db_session: Session
-) -> InputPrompt:
-    query = select(InputPrompt).where(InputPrompt.id == id)
-
-    if user_id:
-        query = query.where(
-            (InputPrompt.user_id == user_id) | (InputPrompt.user_id is None)
-        )
-    else:
-        # If no user_id is provided, only fetch prompts without a user_id (aka public)
-        query = query.where(InputPrompt.user_id == None)  # noqa
-
-    result = db_session.scalar(query)
-
-    if result is None:
-        raise HTTPException(422, "No input prompt found")
-
-    return result
-
-
-def fetch_public_input_prompts(
-    db_session: Session,
-) -> list[InputPrompt]:
-    query = select(InputPrompt).where(InputPrompt.is_public)
-    return list(db_session.scalars(query).all())
-
-
-def fetch_input_prompts_by_user(
-    db_session: Session,
-    user_id: UUID | None,
-    active: bool | None = None,
-    include_public: bool = False,
-) -> list[InputPrompt]:
-    query = select(InputPrompt)
-
-    if user_id is not None:
-        if include_public:
-            query = query.where(
-                (InputPrompt.user_id == user_id) | InputPrompt.is_public
-            )
-        else:
-            query = query.where(InputPrompt.user_id == user_id)
-
-    elif include_public:
-        query = query.where(InputPrompt.is_public)
-
-    if active is not None:
-        query = query.where(InputPrompt.active == active)
-
-    return list(db_session.scalars(query).all())

backend/danswer/db/models.py

@@ -159,9 +159,6 @@ class User(SQLAlchemyBaseUserTableUUID, Base):
     )
 
     prompts: Mapped[list["Prompt"]] = relationship("Prompt", back_populates="user")
-    input_prompts: Mapped[list["InputPrompt"]] = relationship(
-        "InputPrompt", back_populates="user"
-    )
 
     # Personas owned by this user
     personas: Mapped[list["Persona"]] = relationship("Persona", back_populates="user")
@@ -178,31 +175,6 @@ class User(SQLAlchemyBaseUserTableUUID, Base):
     )
 
 
-class InputPrompt(Base):
-    __tablename__ = "inputprompt"
-
-    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
-    prompt: Mapped[str] = mapped_column(String)
-    content: Mapped[str] = mapped_column(String)
-    active: Mapped[bool] = mapped_column(Boolean)
-    user: Mapped[User | None] = relationship("User", back_populates="input_prompts")
-    is_public: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)
-    user_id: Mapped[UUID | None] = mapped_column(
-        ForeignKey("user.id", ondelete="CASCADE"), nullable=True
-    )
-
-
-class InputPrompt__User(Base):
-    __tablename__ = "inputprompt__user"
-
-    input_prompt_id: Mapped[int] = mapped_column(
-        ForeignKey("inputprompt.id"), primary_key=True
-    )
-    user_id: Mapped[UUID | None] = mapped_column(
-        ForeignKey("inputprompt.id"), primary_key=True
-    )
-
-
 class AccessToken(SQLAlchemyBaseAccessTokenTableUUID, Base):
     pass
 
@@ -596,6 +568,25 @@ class Connector(Base):
         list["DocumentByConnectorCredentialPair"]
     ] = relationship("DocumentByConnectorCredentialPair", back_populates="connector")
 
+    # synchronize this validation logic with RefreshFrequencySchema etc on front end
+    # until we have a centralized validation schema
+
+    # TODO(rkuo): experiment with SQLAlchemy validators rather than manual checks
+    # https://docs.sqlalchemy.org/en/20/orm/mapped_attributes.html
+    def validate_refresh_freq(self) -> None:
+        if self.refresh_freq is not None:
+            if self.refresh_freq < 60:
+                raise ValueError(
+                    "refresh_freq must be greater than or equal to 60 seconds."
+                )
+
+    def validate_prune_freq(self) -> None:
+        if self.prune_freq is not None:
+            if self.prune_freq < 86400:
+                raise ValueError(
+                    "prune_freq must be greater than or equal to 86400 seconds."
+                )
+
 
 class Credential(Base):
     __tablename__ = "credential"
@@ -1490,7 +1481,9 @@ class SlackChannelConfig(Base):
     __tablename__ = "slack_channel_config"
 
     id: Mapped[int] = mapped_column(primary_key=True)
-    slack_bot_id: Mapped[int] = mapped_column(ForeignKey("slack_bot.id"), nullable=True)
+    slack_bot_id: Mapped[int] = mapped_column(
+        ForeignKey("slack_bot.id"), nullable=False
+    )
     persona_id: Mapped[int | None] = mapped_column(
         ForeignKey("persona.id"), nullable=True
     )

backend/danswer/db/persona.py

@@ -453,9 +453,9 @@ def upsert_persona(
     """
 
     if persona_id is not None:
-        persona = db_session.query(Persona).filter_by(id=persona_id).first()
+        existing_persona = db_session.query(Persona).filter_by(id=persona_id).first()
     else:
-        persona = _get_persona_by_name(
+        existing_persona = _get_persona_by_name(
             persona_name=name, user=user, db_session=db_session
         )
 
@@ -481,62 +481,78 @@ def upsert_persona(
     prompts = None
     if prompt_ids is not None:
         prompts = db_session.query(Prompt).filter(Prompt.id.in_(prompt_ids)).all()
-        if not prompts and prompt_ids:
-            raise ValueError("prompts not found")
+
+    if prompts is not None and len(prompts) == 0:
+        raise ValueError(
+            f"Invalid Persona config, no valid prompts "
+            f"specified. Specified IDs were: '{prompt_ids}'"
+        )
 
     # ensure all specified tools are valid
     if tools:
         validate_persona_tools(tools)
 
-    if persona:
+    if existing_persona:
         # Built-in personas can only be updated through YAML configuration.
         # This ensures that core system personas are not modified unintentionally.
-        if persona.builtin_persona and not builtin_persona:
+        if existing_persona.builtin_persona and not builtin_persona:
             raise ValueError("Cannot update builtin persona with non-builtin.")
 
         # this checks if the user has permission to edit the persona
-        persona = fetch_persona_by_id(
-            db_session=db_session, persona_id=persona.id, user=user, get_editable=True
+        # will raise an Exception if the user does not have permission
+        existing_persona = fetch_persona_by_id(
+            db_session=db_session,
+            persona_id=existing_persona.id,
+            user=user,
+            get_editable=True,
         )
 
         # The following update excludes `default`, `built-in`, and display priority.
         # Display priority is handled separately in the `display-priority` endpoint.
         # `default` and `built-in` properties can only be set when creating a persona.
-        persona.name = name
-        persona.description = description
-        persona.num_chunks = num_chunks
-        persona.chunks_above = chunks_above
-        persona.chunks_below = chunks_below
-        persona.llm_relevance_filter = llm_relevance_filter
-        persona.llm_filter_extraction = llm_filter_extraction
-        persona.recency_bias = recency_bias
-        persona.llm_model_provider_override = llm_model_provider_override
-        persona.llm_model_version_override = llm_model_version_override
-        persona.starter_messages = starter_messages
-        persona.deleted = False  # Un-delete if previously deleted
-        persona.is_public = is_public
-        persona.icon_color = icon_color
-        persona.icon_shape = icon_shape
+        existing_persona.name = name
+        existing_persona.description = description
+        existing_persona.num_chunks = num_chunks
+        existing_persona.chunks_above = chunks_above
+        existing_persona.chunks_below = chunks_below
+        existing_persona.llm_relevance_filter = llm_relevance_filter
+        existing_persona.llm_filter_extraction = llm_filter_extraction
+        existing_persona.recency_bias = recency_bias
+        existing_persona.llm_model_provider_override = llm_model_provider_override
+        existing_persona.llm_model_version_override = llm_model_version_override
+        existing_persona.starter_messages = starter_messages
+        existing_persona.deleted = False  # Un-delete if previously deleted
+        existing_persona.is_public = is_public
+        existing_persona.icon_color = icon_color
+        existing_persona.icon_shape = icon_shape
         if remove_image or uploaded_image_id:
-            persona.uploaded_image_id = uploaded_image_id
-        persona.is_visible = is_visible
-        persona.search_start_date = search_start_date
-        persona.category_id = category_id
+            existing_persona.uploaded_image_id = uploaded_image_id
+        existing_persona.is_visible = is_visible
+        existing_persona.search_start_date = search_start_date
+        existing_persona.category_id = category_id
         # Do not delete any associations manually added unless
         # a new updated list is provided
         if document_sets is not None:
-            persona.document_sets.clear()
-            persona.document_sets = document_sets or []
+            existing_persona.document_sets.clear()
+            existing_persona.document_sets = document_sets or []
 
         if prompts is not None:
-            persona.prompts.clear()
-            persona.prompts = prompts or []
+            existing_persona.prompts.clear()
+            existing_persona.prompts = prompts
 
         if tools is not None:
-            persona.tools = tools or []
+            existing_persona.tools = tools or []
+
+        persona = existing_persona
 
     else:
-        persona = Persona(
+        if not prompts:
+            raise ValueError(
+                "Invalid Persona config. "
+                "Must specify at least one prompt for a new persona."
+            )
+
+        new_persona = Persona(
             id=persona_id,
             user_id=user.id if user else None,
             is_public=is_public,
@@ -549,7 +565,7 @@ def upsert_persona(
             llm_filter_extraction=llm_filter_extraction,
             recency_bias=recency_bias,
             builtin_persona=builtin_persona,
-            prompts=prompts or [],
+            prompts=prompts,
             document_sets=document_sets or [],
             llm_model_provider_override=llm_model_provider_override,
             llm_model_version_override=llm_model_version_override,
@@ -564,8 +580,8 @@ def upsert_persona(
             is_default_persona=is_default_persona,
             category_id=category_id,
         )
-        db_session.add(persona)
-
+        db_session.add(new_persona)
+        persona = new_persona
     if commit:
         db_session.commit()
     else:

backend/danswer/document_index/vespa/app_config/schemas/danswer_chunk.sd

@@ -4,6 +4,8 @@ schema DANSWER_CHUNK_NAME {
         # Not to be confused with the UUID generated for this chunk which is called documentid by default
         field document_id type string {
             indexing: summary | attribute
+            attribute: fast-search
+            rank: filter
         }
         field chunk_id type int {
             indexing: summary | attribute

backend/danswer/file_processing/extract_file_text.py

@@ -6,6 +6,7 @@ import zipfile
 from collections.abc import Callable
 from collections.abc import Iterator
 from email.parser import Parser as EmailParser
+from io import BytesIO
 from pathlib import Path
 from typing import Any
 from typing import Dict
@@ -15,13 +16,17 @@ import chardet
 import docx  # type: ignore
 import openpyxl  # type: ignore
 import pptx  # type: ignore
+from docx import Document
+from fastapi import UploadFile
 from pypdf import PdfReader
 from pypdf.errors import PdfStreamError
 
 from danswer.configs.constants import DANSWER_METADATA_FILENAME
+from danswer.configs.constants import FileOrigin
 from danswer.file_processing.html_utils import parse_html_page_basic
 from danswer.file_processing.unstructured import get_unstructured_api_key
 from danswer.file_processing.unstructured import unstructured_to_text
+from danswer.file_store.file_store import FileStore
 from danswer.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -65,7 +70,7 @@ def get_file_ext(file_path_or_name: str | Path) -> str:
     return extension
 
 
-def check_file_ext_is_valid(ext: str) -> bool:
+def is_valid_file_ext(ext: str) -> bool:
     return ext in VALID_FILE_EXTENSIONS
 
 
@@ -359,7 +364,7 @@ def extract_file_text(
             elif file_name is not None:
                 final_extension = get_file_ext(file_name)
 
-            if check_file_ext_is_valid(final_extension):
+            if is_valid_file_ext(final_extension):
                 return extension_to_function.get(final_extension, file_io_to_text)(file)
 
         # Either the file somehow has no name or the extension is not one that we recognize
@@ -375,3 +380,35 @@ def extract_file_text(
             ) from e
         logger.warning(f"Failed to process file {file_name or 'Unknown'}: {str(e)}")
         return ""
+
+
+def convert_docx_to_txt(
+    file: UploadFile, file_store: FileStore, file_path: str
+) -> None:
+    file.file.seek(0)
+    docx_content = file.file.read()
+    doc = Document(BytesIO(docx_content))
+
+    # Extract text from the document
+    full_text = []
+    for para in doc.paragraphs:
+        full_text.append(para.text)
+
+    # Join the extracted text
+    text_content = "\n".join(full_text)
+
+    txt_file_path = docx_to_txt_filename(file_path)
+    file_store.save_file(
+        file_name=txt_file_path,
+        content=BytesIO(text_content.encode("utf-8")),
+        display_name=file.filename,
+        file_origin=FileOrigin.CONNECTOR,
+        file_type="text/plain",
+    )
+
+
+def docx_to_txt_filename(file_path: str) -> str:
+    """
+    Convert a .docx file path to its corresponding .txt file path.
+    """
+    return file_path.rsplit(".", 1)[0] + ".txt"

backend/danswer/file_store/utils.py

@@ -1,6 +1,6 @@
+import base64
 from collections.abc import Callable
 from io import BytesIO
-from typing import Any
 from typing import cast
 from uuid import uuid4
 
@@ -13,8 +13,8 @@ from danswer.db.models import ChatMessage
 from danswer.file_store.file_store import get_default_file_store
 from danswer.file_store.models import FileDescriptor
 from danswer.file_store.models import InMemoryChatFile
+from danswer.utils.b64 import get_image_type
 from danswer.utils.threadpool_concurrency import run_functions_tuples_in_parallel
-from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 
 
 def load_chat_file(
@@ -75,11 +75,58 @@ def save_file_from_url(url: str, tenant_id: str) -> str:
         return unique_id
 
 
-def save_files_from_urls(urls: list[str]) -> list[str]:
-    tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
+def save_file_from_base64(base64_string: str, tenant_id: str) -> str:
+    with get_session_with_tenant(tenant_id) as db_session:
+        unique_id = str(uuid4())
+        file_store = get_default_file_store(db_session)
+        file_store.save_file(
+            file_name=unique_id,
+            content=BytesIO(base64.b64decode(base64_string)),
+            display_name="GeneratedImage",
+            file_origin=FileOrigin.CHAT_IMAGE_GEN,
+            file_type=get_image_type(base64_string),
+        )
+        return unique_id
 
-    funcs: list[tuple[Callable[..., Any], tuple[Any, ...]]] = [
-        (save_file_from_url, (url, tenant_id)) for url in urls
+
+def save_file(
+    tenant_id: str,
+    url: str | None = None,
+    base64_data: str | None = None,
+) -> str:
+    """Save a file from either a URL or base64 encoded string.
+
+    Args:
+        tenant_id: The tenant ID to save the file under
+        url: URL to download file from
+        base64_data: Base64 encoded file data
+
+    Returns:
+        The unique ID of the saved file
+
+    Raises:
+        ValueError: If neither url nor base64_data is provided, or if both are provided
+    """
+    if url is not None and base64_data is not None:
+        raise ValueError("Cannot specify both url and base64_data")
+
+    if url is not None:
+        return save_file_from_url(url, tenant_id)
+    elif base64_data is not None:
+        return save_file_from_base64(base64_data, tenant_id)
+    else:
+        raise ValueError("Must specify either url or base64_data")
+
+
+def save_files(urls: list[str], base64_files: list[str], tenant_id: str) -> list[str]:
+    # NOTE: be explicit about typing so that if we change things, we get notified
+    funcs: list[
+        tuple[
+            Callable[[str, str | None, str | None], str],
+            tuple[str, str | None, str | None],
+        ]
+    ] = [(save_file, (tenant_id, url, None)) for url in urls] + [
+        (save_file, (tenant_id, None, base64_file)) for base64_file in base64_files
     ]
-    # Must pass in tenant_id here, since this is called by multithreading
+
     return run_functions_tuples_in_parallel(funcs)

backend/danswer/indexing/chunker.py

@@ -375,6 +375,6 @@ class Chunker:
             final_chunks.extend(chunks)
 
             if self.callback:
-                self.callback.progress("Chunker.chunk", len(chunks))
+                self.callback.progress("Chunker.chunk", 0)
 
         return final_chunks

backend/danswer/indexing/indexing_heartbeat.py

@@ -12,4 +12,5 @@ class IndexingHeartbeatInterface(ABC):
 
     @abstractmethod
     def progress(self, tag: str, amount: int) -> None:
-        """Send progress updates to the caller."""
+        """Send progress updates to the caller.
+        Send 0 for the amount if you just want to do a keep alive ping."""

backend/danswer/indexing/indexing_pipeline.py

@@ -1,4 +1,5 @@
 import traceback
+from collections.abc import Callable
 from functools import partial
 from http import HTTPStatus
 from typing import Protocol
@@ -12,6 +13,7 @@ from danswer.access.access import get_access_for_documents
 from danswer.access.models import DocumentAccess
 from danswer.configs.app_configs import ENABLE_MULTIPASS_INDEXING
 from danswer.configs.app_configs import INDEXING_EXCEPTION_LIMIT
+from danswer.configs.app_configs import MAX_DOCUMENT_CHARS
 from danswer.configs.constants import DEFAULT_BOOST
 from danswer.connectors.cross_connector_utils.miscellaneous_utils import (
     get_experts_stores_representations,
@@ -202,40 +204,13 @@ def index_doc_batch_with_handler(
 
 
 def index_doc_batch_prepare(
-    document_batch: list[Document],
+    documents: list[Document],
     index_attempt_metadata: IndexAttemptMetadata,
     db_session: Session,
     ignore_time_skip: bool = False,
 ) -> DocumentBatchPrepareContext | None:
     """Sets up the documents in the relational DB (source of truth) for permissions, metadata, etc.
     This preceeds indexing it into the actual document index."""
-    documents: list[Document] = []
-    for document in document_batch:
-        empty_contents = not any(section.text.strip() for section in document.sections)
-        if (
-            (not document.title or not document.title.strip())
-            and not document.semantic_identifier.strip()
-            and empty_contents
-        ):
-            # Skip documents that have neither title nor content
-            # If the document doesn't have either, then there is no useful information in it
-            # This is again verified later in the pipeline after chunking but at that point there should
-            # already be no documents that are empty.
-            logger.warning(
-                f"Skipping document with ID {document.id} as it has neither title nor content."
-            )
-            continue
-
-        if document.title is not None and not document.title.strip() and empty_contents:
-            # The title is explicitly empty ("" and not None) and the document is empty
-            # so when building the chunk text representation, it will be empty and unuseable
-            logger.warning(
-                f"Skipping document with ID {document.id} as the chunks will be empty."
-            )
-            continue
-
-        documents.append(document)
-
     # Create a trimmed list of docs that don't have a newer updated at
     # Shortcuts the time-consuming flow on connector index retries
     document_ids: list[str] = [document.id for document in documents]
@@ -282,17 +257,64 @@ def index_doc_batch_prepare(
     )
 
 
+def filter_documents(document_batch: list[Document]) -> list[Document]:
+    documents: list[Document] = []
+    for document in document_batch:
+        empty_contents = not any(section.text.strip() for section in document.sections)
+        if (
+            (not document.title or not document.title.strip())
+            and not document.semantic_identifier.strip()
+            and empty_contents
+        ):
+            # Skip documents that have neither title nor content
+            # If the document doesn't have either, then there is no useful information in it
+            # This is again verified later in the pipeline after chunking but at that point there should
+            # already be no documents that are empty.
+            logger.warning(
+                f"Skipping document with ID {document.id} as it has neither title nor content."
+            )
+            continue
+
+        if document.title is not None and not document.title.strip() and empty_contents:
+            # The title is explicitly empty ("" and not None) and the document is empty
+            # so when building the chunk text representation, it will be empty and unuseable
+            logger.warning(
+                f"Skipping document with ID {document.id} as the chunks will be empty."
+            )
+            continue
+
+        section_chars = sum(len(section.text) for section in document.sections)
+        if (
+            MAX_DOCUMENT_CHARS
+            and len(document.title or document.semantic_identifier) + section_chars
+            > MAX_DOCUMENT_CHARS
+        ):
+            # Skip documents that are too long, later on there are more memory intensive steps done on the text
+            # and the container will run out of memory and crash. Several other checks are included upstream but
+            # those are at the connector level so a catchall is still needed.
+            # Assumption here is that files that are that long, are generated files and not the type users
+            # generally care for.
+            logger.warning(
+                f"Skipping document with ID {document.id} as it is too long."
+            )
+            continue
+
+        documents.append(document)
+    return documents
+
+
 @log_function_time(debug_only=True)
 def index_doc_batch(
     *,
+    document_batch: list[Document],
     chunker: Chunker,
     embedder: IndexingEmbedder,
     document_index: DocumentIndex,
-    document_batch: list[Document],
     index_attempt_metadata: IndexAttemptMetadata,
     db_session: Session,
     ignore_time_skip: bool = False,
     tenant_id: str | None = None,
+    filter_fnc: Callable[[list[Document]], list[Document]] = filter_documents,
 ) -> tuple[int, int]:
     """Takes different pieces of the indexing pipeline and applies it to a batch of documents
     Note that the documents should already be batched at this point so that it does not inflate the
@@ -309,8 +331,11 @@ def index_doc_batch(
         is_public=False,
     )
 
+    logger.debug("Filtering Documents")
+    filtered_documents = filter_fnc(document_batch)
+
     ctx = index_doc_batch_prepare(
-        document_batch=document_batch,
+        documents=filtered_documents,
         index_attempt_metadata=index_attempt_metadata,
         ignore_time_skip=ignore_time_skip,
         db_session=db_session,

backend/danswer/llm/answering/models.py

@@ -1,163 +0,0 @@
-from collections.abc import Callable
-from collections.abc import Iterator
-from typing import TYPE_CHECKING
-
-from langchain.schema.messages import AIMessage
-from langchain.schema.messages import BaseMessage
-from langchain.schema.messages import HumanMessage
-from langchain.schema.messages import SystemMessage
-from pydantic import BaseModel
-from pydantic import ConfigDict
-from pydantic import Field
-from pydantic import model_validator
-
-from danswer.chat.models import AnswerQuestionStreamReturn
-from danswer.configs.constants import MessageType
-from danswer.file_store.models import InMemoryChatFile
-from danswer.llm.override_models import PromptOverride
-from danswer.llm.utils import build_content_with_imgs
-from danswer.tools.models import ToolCallFinalResult
-
-if TYPE_CHECKING:
-    from danswer.db.models import ChatMessage
-    from danswer.db.models import Prompt
-
-
-StreamProcessor = Callable[[Iterator[str]], AnswerQuestionStreamReturn]
-
-
-class PreviousMessage(BaseModel):
-    """Simplified version of `ChatMessage`"""
-
-    message: str
-    token_count: int
-    message_type: MessageType
-    files: list[InMemoryChatFile]
-    tool_call: ToolCallFinalResult | None
-
-    @classmethod
-    def from_chat_message(
-        cls, chat_message: "ChatMessage", available_files: list[InMemoryChatFile]
-    ) -> "PreviousMessage":
-        message_file_ids = (
-            [file["id"] for file in chat_message.files] if chat_message.files else []
-        )
-        return cls(
-            message=chat_message.message,
-            token_count=chat_message.token_count,
-            message_type=chat_message.message_type,
-            files=[
-                file
-                for file in available_files
-                if str(file.file_id) in message_file_ids
-            ],
-            tool_call=ToolCallFinalResult(
-                tool_name=chat_message.tool_call.tool_name,
-                tool_args=chat_message.tool_call.tool_arguments,
-                tool_result=chat_message.tool_call.tool_result,
-            )
-            if chat_message.tool_call
-            else None,
-        )
-
-    def to_langchain_msg(self) -> BaseMessage:
-        content = build_content_with_imgs(self.message, self.files)
-        if self.message_type == MessageType.USER:
-            return HumanMessage(content=content)
-        elif self.message_type == MessageType.ASSISTANT:
-            return AIMessage(content=content)
-        else:
-            return SystemMessage(content=content)
-
-
-class DocumentPruningConfig(BaseModel):
-    max_chunks: int | None = None
-    max_window_percentage: float | None = None
-    max_tokens: int | None = None
-    # different pruning behavior is expected when the
-    # user manually selects documents they want to chat with
-    # e.g. we don't want to truncate each document to be no more
-    # than one chunk long
-    is_manually_selected_docs: bool = False
-    # If user specifies to include additional context Chunks for each match, then different pruning
-    # is used. As many Sections as possible are included, and the last Section is truncated
-    # If this is false, all of the Sections are truncated if they are longer than the expected Chunk size.
-    # Sections are often expected to be longer than the maximum Chunk size but Chunks should not be.
-    use_sections: bool = True
-    # If using tools, then we need to consider the tool length
-    tool_num_tokens: int = 0
-    # If using a tool message to represent the docs, then we have to JSON serialize
-    # the document content, which adds to the token count.
-    using_tool_message: bool = False
-
-
-class ContextualPruningConfig(DocumentPruningConfig):
-    num_chunk_multiple: int
-
-    @classmethod
-    def from_doc_pruning_config(
-        cls, num_chunk_multiple: int, doc_pruning_config: DocumentPruningConfig
-    ) -> "ContextualPruningConfig":
-        return cls(num_chunk_multiple=num_chunk_multiple, **doc_pruning_config.dict())
-
-
-class CitationConfig(BaseModel):
-    all_docs_useful: bool = False
-
-
-class QuotesConfig(BaseModel):
-    pass
-
-
-class AnswerStyleConfig(BaseModel):
-    citation_config: CitationConfig | None = None
-    quotes_config: QuotesConfig | None = None
-    document_pruning_config: DocumentPruningConfig = Field(
-        default_factory=DocumentPruningConfig
-    )
-    # forces the LLM to return a structured response, see
-    # https://platform.openai.com/docs/guides/structured-outputs/introduction
-    # right now, only used by the simple chat API
-    structured_response_format: dict | None = None
-
-    @model_validator(mode="after")
-    def check_quotes_and_citation(self) -> "AnswerStyleConfig":
-        if self.citation_config is None and self.quotes_config is None:
-            raise ValueError(
-                "One of `citation_config` or `quotes_config` must be provided"
-            )
-
-        if self.citation_config is not None and self.quotes_config is not None:
-            raise ValueError(
-                "Only one of `citation_config` or `quotes_config` must be provided"
-            )
-
-        return self
-
-
-class PromptConfig(BaseModel):
-    """Final representation of the Prompt configuration passed
-    into the `Answer` object."""
-
-    system_prompt: str
-    task_prompt: str
-    datetime_aware: bool
-    include_citations: bool
-
-    @classmethod
-    def from_model(
-        cls, model: "Prompt", prompt_override: PromptOverride | None = None
-    ) -> "PromptConfig":
-        override_system_prompt = (
-            prompt_override.system_prompt if prompt_override else None
-        )
-        override_task_prompt = prompt_override.task_prompt if prompt_override else None
-
-        return cls(
-            system_prompt=override_system_prompt or model.system_prompt,
-            task_prompt=override_task_prompt or model.task_prompt,
-            datetime_aware=model.datetime_aware,
-            include_citations=model.include_citations,
-        )
-
-    model_config = ConfigDict(frozen=True)

backend/danswer/llm/answering/prompts/utils.py

@@ -1,20 +0,0 @@
-from danswer.prompts.direct_qa_prompts import PARAMATERIZED_PROMPT
-from danswer.prompts.direct_qa_prompts import PARAMATERIZED_PROMPT_WITHOUT_CONTEXT
-
-
-def build_dummy_prompt(
-    system_prompt: str, task_prompt: str, retrieval_disabled: bool
-) -> str:
-    if retrieval_disabled:
-        return PARAMATERIZED_PROMPT_WITHOUT_CONTEXT.format(
-            user_query="<USER_QUERY>",
-            system_prompt=system_prompt,
-            task_prompt=task_prompt,
-        ).strip()
-
-    return PARAMATERIZED_PROMPT.format(
-        context_docs_str="<CONTEXT_DOCS>",
-        user_query="<USER_QUERY>",
-        system_prompt=system_prompt,
-        task_prompt=task_prompt,
-    ).strip()

backend/danswer/llm/chat_llm.py

@@ -268,12 +268,16 @@ class DefaultMultiLLM(LLM):
 
         # NOTE: have to set these as environment variables for Litellm since
         # not all are able to passed in but they always support them set as env
-        # variables
+        # variables. We'll also try passing them in, since litellm just ignores
+        # addtional kwargs (and some kwargs MUST be passed in rather than set as
+        # env variables)
         if custom_config:
             for k, v in custom_config.items():
                 os.environ[k] = v
 
         model_kwargs = model_kwargs or {}
+        if custom_config:
+            model_kwargs.update(custom_config)
         if extra_headers:
             model_kwargs.update({"extra_headers": extra_headers})
         if extra_body:

backend/danswer/llm/models.py

@@ -0,0 +1,59 @@
+from typing import TYPE_CHECKING
+
+from langchain.schema.messages import AIMessage
+from langchain.schema.messages import BaseMessage
+from langchain.schema.messages import HumanMessage
+from langchain.schema.messages import SystemMessage
+from pydantic import BaseModel
+
+from danswer.configs.constants import MessageType
+from danswer.file_store.models import InMemoryChatFile
+from danswer.llm.utils import build_content_with_imgs
+from danswer.tools.models import ToolCallFinalResult
+
+if TYPE_CHECKING:
+    from danswer.db.models import ChatMessage
+
+
+class PreviousMessage(BaseModel):
+    """Simplified version of `ChatMessage`"""
+
+    message: str
+    token_count: int
+    message_type: MessageType
+    files: list[InMemoryChatFile]
+    tool_call: ToolCallFinalResult | None
+
+    @classmethod
+    def from_chat_message(
+        cls, chat_message: "ChatMessage", available_files: list[InMemoryChatFile]
+    ) -> "PreviousMessage":
+        message_file_ids = (
+            [file["id"] for file in chat_message.files] if chat_message.files else []
+        )
+        return cls(
+            message=chat_message.message,
+            token_count=chat_message.token_count,
+            message_type=chat_message.message_type,
+            files=[
+                file
+                for file in available_files
+                if str(file.file_id) in message_file_ids
+            ],
+            tool_call=ToolCallFinalResult(
+                tool_name=chat_message.tool_call.tool_name,
+                tool_args=chat_message.tool_call.tool_arguments,
+                tool_result=chat_message.tool_call.tool_result,
+            )
+            if chat_message.tool_call
+            else None,
+        )
+
+    def to_langchain_msg(self) -> BaseMessage:
+        content = build_content_with_imgs(self.message, self.files)
+        if self.message_type == MessageType.USER:
+            return HumanMessage(content=content)
+        elif self.message_type == MessageType.ASSISTANT:
+            return AIMessage(content=content)
+        else:
+            return SystemMessage(content=content)

backend/danswer/llm/utils.py

@@ -1,15 +1,11 @@
 import copy
-import io
 import json
 from collections.abc import Callable
 from collections.abc import Iterator
 from typing import Any
 from typing import cast
-from typing import TYPE_CHECKING
-from typing import Union
 
 import litellm  # type: ignore
-import pandas as pd
 import tiktoken
 from langchain.prompts.base import StringPromptValue
 from langchain.prompts.chat import ChatPromptValue
@@ -36,17 +32,15 @@ from danswer.configs.constants import MessageType
 from danswer.configs.model_configs import GEN_AI_MAX_TOKENS
 from danswer.configs.model_configs import GEN_AI_MODEL_FALLBACK_MAX_TOKENS
 from danswer.configs.model_configs import GEN_AI_NUM_RESERVED_OUTPUT_TOKENS
-from danswer.db.models import ChatMessage
 from danswer.file_store.models import ChatFileType
 from danswer.file_store.models import InMemoryChatFile
 from danswer.llm.interfaces import LLM
 from danswer.prompts.constants import CODE_BLOCK_PAT
+from danswer.utils.b64 import get_image_type
+from danswer.utils.b64 import get_image_type_from_bytes
 from danswer.utils.logger import setup_logger
 from shared_configs.configs import LOG_LEVEL
 
-if TYPE_CHECKING:
-    from danswer.llm.answering.models import PreviousMessage
-
 logger = setup_logger()
 
 
@@ -104,92 +98,39 @@ def litellm_exception_to_error_msg(
     return error_msg
 
 
-def translate_danswer_msg_to_langchain(
-    msg: Union[ChatMessage, "PreviousMessage"],
-) -> BaseMessage:
-    files: list[InMemoryChatFile] = []
-
-    # If the message is a `ChatMessage`, it doesn't have the downloaded files
-    # attached. Just ignore them for now.
-    if not isinstance(msg, ChatMessage):
-        files = msg.files
-    content = build_content_with_imgs(msg.message, files, message_type=msg.message_type)
-
-    if msg.message_type == MessageType.SYSTEM:
-        raise ValueError("System messages are not currently part of history")
-    if msg.message_type == MessageType.ASSISTANT:
-        return AIMessage(content=content)
-    if msg.message_type == MessageType.USER:
-        return HumanMessage(content=content)
-
-    raise ValueError(f"New message type {msg.message_type} not handled")
-
-
-def translate_history_to_basemessages(
-    history: list[ChatMessage] | list["PreviousMessage"],
-) -> tuple[list[BaseMessage], list[int]]:
-    history_basemessages = [
-        translate_danswer_msg_to_langchain(msg)
-        for msg in history
-        if msg.token_count != 0
-    ]
-    history_token_counts = [msg.token_count for msg in history if msg.token_count != 0]
-    return history_basemessages, history_token_counts
-
-
-#  Processes CSV files to show the first 5 rows and max_columns (default 40) columns
-def _process_csv_file(file: InMemoryChatFile, max_columns: int = 40) -> str:
-    df = pd.read_csv(io.StringIO(file.content.decode("utf-8")))
-
-    csv_preview = df.head().to_string(max_cols=max_columns)
-
-    file_name_section = (
-        f"CSV FILE NAME: {file.filename}\n"
-        if file.filename
-        else "CSV FILE (NO NAME PROVIDED):\n"
-    )
-    return f"{file_name_section}{CODE_BLOCK_PAT.format(csv_preview)}\n\n\n"
-
-
 def _build_content(
     message: str,
     files: list[InMemoryChatFile] | None = None,
 ) -> str:
     """Applies all non-image files."""
-    text_files = (
-        [file for file in files if file.file_type == ChatFileType.PLAIN_TEXT]
-        if files
-        else None
-    )
+    if not files:
+        return message
 
-    csv_files = (
-        [file for file in files if file.file_type == ChatFileType.CSV]
-        if files
-        else None
-    )
+    text_files = [
+        file
+        for file in files
+        if file.file_type in (ChatFileType.PLAIN_TEXT, ChatFileType.CSV)
+    ]
 
-    if not text_files and not csv_files:
+    if not text_files:
         return message
 
     final_message_with_files = "FILES:\n\n"
-    for file in text_files or []:
+    for file in text_files:
         file_content = file.content.decode("utf-8")
         file_name_section = f"DOCUMENT: {file.filename}\n" if file.filename else ""
         final_message_with_files += (
             f"{file_name_section}{CODE_BLOCK_PAT.format(file_content.strip())}\n\n\n"
         )
-    for file in csv_files or []:
-        final_message_with_files += _process_csv_file(file)
 
-    final_message_with_files += message
-
-    return final_message_with_files
+    return final_message_with_files + message
 
 
 def build_content_with_imgs(
     message: str,
     files: list[InMemoryChatFile] | None = None,
     img_urls: list[str] | None = None,
+    b64_imgs: list[str] | None = None,
     message_type: MessageType = MessageType.USER,
 ) -> str | list[str | dict[str, Any]]:  # matching Langchain's BaseMessage content type
     files = files or []
@@ -202,6 +143,7 @@ def build_content_with_imgs(
     )
 
     img_urls = img_urls or []
+    b64_imgs = b64_imgs or []
 
     message_main_content = _build_content(message, files)
 
@@ -220,11 +162,22 @@ def build_content_with_imgs(
             {
                 "type": "image_url",
                 "image_url": {
-                    "url": f"data:image/jpeg;base64,{file.to_base64()}",
+                    "url": (
+                        f"data:{get_image_type_from_bytes(file.content)};"
+                        f"base64,{file.to_base64()}"
+                    ),
                 },
             }
-            for file in files
-            if file.file_type == "image"
+            for file in img_files
+        ]
+        + [
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": f"data:{get_image_type(b64_img)};base64,{b64_img}",
+                },
+            }
+            for b64_img in b64_imgs
         ]
         + [
             {

backend/danswer/main.py

@@ -52,12 +52,9 @@ from danswer.server.documents.connector import router as connector_router
 from danswer.server.documents.credential import router as credential_router
 from danswer.server.documents.document import router as document_router
 from danswer.server.documents.indexing import router as indexing_router
+from danswer.server.documents.standard_oauth import router as oauth_router
 from danswer.server.features.document_set.api import router as document_set_router
 from danswer.server.features.folder.api import router as folder_router
-from danswer.server.features.input_prompt.api import (
-    admin_router as admin_input_prompt_router,
-)
-from danswer.server.features.input_prompt.api import basic_router as input_prompt_router
 from danswer.server.features.notifications.api import router as notification_router
 from danswer.server.features.persona.api import admin_router as admin_persona_router
 from danswer.server.features.persona.api import basic_router as persona_router
@@ -105,7 +102,6 @@ from shared_configs.configs import CORS_ALLOWED_ORIGIN
 from shared_configs.configs import MULTI_TENANT
 from shared_configs.configs import SENTRY_DSN
 
-
 logger = setup_logger()
 
 
@@ -259,8 +255,6 @@ def get_application() -> FastAPI:
     )
     include_router_with_global_prefix_prepended(application, persona_router)
     include_router_with_global_prefix_prepended(application, admin_persona_router)
-    include_router_with_global_prefix_prepended(application, input_prompt_router)
-    include_router_with_global_prefix_prepended(application, admin_input_prompt_router)
     include_router_with_global_prefix_prepended(application, notification_router)
     include_router_with_global_prefix_prepended(application, prompt_router)
     include_router_with_global_prefix_prepended(application, tool_router)
@@ -283,6 +277,7 @@ def get_application() -> FastAPI:
     )
     include_router_with_global_prefix_prepended(application, long_term_logs_router)
     include_router_with_global_prefix_prepended(application, api_key_router)
+    include_router_with_global_prefix_prepended(application, oauth_router)
 
     if AUTH_TYPE == AuthType.DISABLED:
         # Server logs this during auth setup verification step

backend/danswer/natural_language_processing/search_nlp_models.py

@@ -182,7 +182,7 @@ class EmbeddingModel:
             embeddings.extend(response.embeddings)
 
             if self.callback:
-                self.callback.progress("_batch_encode_texts", 1)
+                self.callback.progress("_batch_encode_texts", 0)
         return embeddings
 
     def encode(

backend/danswer/prompts/prompt_utils.py

@@ -5,11 +5,11 @@ from typing import cast
 from langchain_core.messages import BaseMessage
 
 from danswer.chat.models import LlmDoc
+from danswer.chat.models import PromptConfig
 from danswer.configs.chat_configs import LANGUAGE_HINT
 from danswer.configs.constants import DocumentSource
 from danswer.context.search.models import InferenceChunk
 from danswer.db.models import Prompt
-from danswer.llm.answering.models import PromptConfig
 from danswer.prompts.chat_prompts import ADDITIONAL_INFO
 from danswer.prompts.chat_prompts import CITATION_REMINDER
 from danswer.prompts.constants import CODE_BLOCK_PAT

backend/danswer/redis/redis_connector.py

@@ -72,24 +72,6 @@ class RedisConnector:
 
         return finished
 
-    @staticmethod
-    def get_id_from_fence_key(key: str) -> str | None:
-        """
-        Extracts the object ID from a fence key in the format `PREFIX_fence_X`.
-
-        Args:
-            key (str): The fence key string.
-
-        Returns:
-            Optional[int]: The extracted ID if the key is in the correct format, otherwise None.
-        """
-        parts = key.split("_")
-        if len(parts) != 3:
-            return None
-
-        object_id = parts[2]
-        return object_id
-
     @staticmethod
     def get_id_from_task_id(task_id: str) -> str | None:
         """

backend/danswer/redis/redis_connector_index.py

@@ -5,6 +5,13 @@ from uuid import uuid4
 import redis
 from pydantic import BaseModel
 
+from danswer.redis.redis_connector_utils import RedisConnectorUtils
+
+
+class RedisConnectorIndexIdentifiers(BaseModel):
+    cc_pair_id: int
+    search_settings_id: int
+
 
 class RedisConnectorIndexPayload(BaseModel):
     index_attempt_id: int | None
@@ -119,7 +126,8 @@ class RedisConnectorIndex:
         self.redis.delete(self.generator_complete_key)
 
     def get_progress(self) -> int | None:
-        """Returns None if the key doesn't exist. The"""
+        """Returns None if the key doesn't exist. Returns an int representing
+        the number of documents processed otherwise."""
         # TODO: move into fence?
         bytes = self.redis.get(self.generator_progress_key)
         if bytes is None:
@@ -143,6 +151,26 @@ class RedisConnectorIndex:
         self.redis.delete(self.generator_complete_key)
         self.redis.delete(self.fence_key)
 
+    @staticmethod
+    def parse_key(key_bytes: bytes) -> RedisConnectorIndexIdentifiers | None:
+        """Parses the redis key for an indexing job and extracts
+        the cc_pair_id and search_settings_id"""
+        fence_key = key_bytes.decode("utf-8")
+        composite_id = RedisConnectorUtils.get_id_from_fence_key(fence_key)
+        if composite_id is None:
+            return None
+
+        # parse out metadata and initialize the helper class with it
+        parts = composite_id.split("/")
+        if len(parts) != 2:
+            return None
+
+        cc_pair_id = int(parts[0])
+        search_settings_id = int(parts[1])
+        return RedisConnectorIndexIdentifiers(
+            cc_pair_id=cc_pair_id, search_settings_id=search_settings_id
+        )
+
     @staticmethod
     def reset_all(r: redis.Redis) -> None:
         """Deletes all redis values for all connectors"""

backend/danswer/redis/redis_connector_utils.py

@@ -0,0 +1,18 @@
+class RedisConnectorUtils:
+    @staticmethod
+    def get_id_from_fence_key(key: str) -> str | None:
+        """
+        Extracts the object ID from a fence key in the format `PREFIX_fence_X`.
+
+        Args:
+            key (str): The fence key string.
+
+        Returns:
+            Optional[int]: The extracted ID if the key is in the correct format, otherwise None.
+        """
+        parts = key.split("_")
+        if len(parts) != 3:
+            return None
+
+        object_id = parts[2]
+        return object_id

backend/danswer/secondary_llm_flows/choose_search.py

@@ -3,14 +3,14 @@ from langchain.schema import HumanMessage
 from langchain.schema import SystemMessage
 
 from danswer.chat.chat_utils import combine_message_chain
+from danswer.chat.prompt_builder.utils import translate_danswer_msg_to_langchain
 from danswer.configs.chat_configs import DISABLE_LLM_CHOOSE_SEARCH
 from danswer.configs.model_configs import GEN_AI_HISTORY_CUTOFF
 from danswer.db.models import ChatMessage
-from danswer.llm.answering.models import PreviousMessage
 from danswer.llm.interfaces import LLM
+from danswer.llm.models import PreviousMessage
 from danswer.llm.utils import dict_based_prompt_to_langchain_prompt
 from danswer.llm.utils import message_to_string
-from danswer.llm.utils import translate_danswer_msg_to_langchain
 from danswer.prompts.chat_prompts import AGGRESSIVE_SEARCH_TEMPLATE
 from danswer.prompts.chat_prompts import NO_SEARCH
 from danswer.prompts.chat_prompts import REQUIRE_SEARCH_HINT