Skip to content

Commit ecf4923

Browse files
Wevespablonyx
andauthored
Fix answer with specified doc ids (#2703)
* Fix Fix Refactor more more fix refactor Fix circular imports Refactor Move tests around * Add quote support * Testing * More testing * Fix image generation slowness * Remove unused exception * Fix UT * fix stop generating * minor typo * minor logging updates for clarity --------- Co-authored-by: pablodanswer <pablo@danswer.ai>
1 parent d66b81a commit ecf4923

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+1987
-1110
lines changed

backend/danswer/chat/models.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from danswer.search.enums import SearchType
1111
from danswer.search.models import RetrievalDocs
1212
from danswer.search.models import SearchResponse
13-
from danswer.tools.custom.base_tool_types import ToolResultType
13+
from danswer.tools.tool_implementations.custom.base_tool_types import ToolResultType
1414

1515

1616
class LlmDoc(BaseModel):

backend/danswer/chat/process_message.py

Lines changed: 61 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from danswer.chat.models import MessageSpecificCitations
1919
from danswer.chat.models import QADocsResponse
2020
from danswer.chat.models import StreamingError
21+
from danswer.chat.models import StreamStopInfo
2122
from danswer.configs.app_configs import AZURE_DALLE_API_BASE
2223
from danswer.configs.app_configs import AZURE_DALLE_API_KEY
2324
from danswer.configs.app_configs import AZURE_DALLE_API_VERSION
@@ -77,31 +78,49 @@
7778
from danswer.server.query_and_chat.models import CreateChatMessageRequest
7879
from danswer.server.utils import get_json_line
7980
from danswer.tools.built_in_tools import get_built_in_tool_by_id
80-
from danswer.tools.custom.custom_tool import (
81+
from danswer.tools.force import ForceUseTool
82+
from danswer.tools.models import DynamicSchemaInfo
83+
from danswer.tools.models import ToolResponse
84+
from danswer.tools.tool import Tool
85+
from danswer.tools.tool_implementations.custom.custom_tool import (
8186
build_custom_tools_from_openapi_schema_and_headers,
8287
)
83-
from danswer.tools.custom.custom_tool import CUSTOM_TOOL_RESPONSE_ID
84-
from danswer.tools.custom.custom_tool import CustomToolCallSummary
85-
from danswer.tools.force import ForceUseTool
86-
from danswer.tools.images.image_generation_tool import IMAGE_GENERATION_RESPONSE_ID
87-
from danswer.tools.images.image_generation_tool import ImageGenerationResponse
88-
from danswer.tools.images.image_generation_tool import ImageGenerationTool
89-
from danswer.tools.internet_search.internet_search_tool import (
88+
from danswer.tools.tool_implementations.custom.custom_tool import (
89+
CUSTOM_TOOL_RESPONSE_ID,
90+
)
91+
from danswer.tools.tool_implementations.custom.custom_tool import CustomToolCallSummary
92+
from danswer.tools.tool_implementations.images.image_generation_tool import (
93+
IMAGE_GENERATION_RESPONSE_ID,
94+
)
95+
from danswer.tools.tool_implementations.images.image_generation_tool import (
96+
ImageGenerationResponse,
97+
)
98+
from danswer.tools.tool_implementations.images.image_generation_tool import (
99+
ImageGenerationTool,
100+
)
101+
from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
90102
INTERNET_SEARCH_RESPONSE_ID,
91103
)
92-
from danswer.tools.internet_search.internet_search_tool import (
104+
from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
93105
internet_search_response_to_search_docs,
94106
)
95-
from danswer.tools.internet_search.internet_search_tool import InternetSearchResponse
96-
from danswer.tools.internet_search.internet_search_tool import InternetSearchTool
97-
from danswer.tools.models import DynamicSchemaInfo
98-
from danswer.tools.search.search_tool import FINAL_CONTEXT_DOCUMENTS_ID
99-
from danswer.tools.search.search_tool import SEARCH_RESPONSE_SUMMARY_ID
100-
from danswer.tools.search.search_tool import SearchResponseSummary
101-
from danswer.tools.search.search_tool import SearchTool
102-
from danswer.tools.search.search_tool import SECTION_RELEVANCE_LIST_ID
103-
from danswer.tools.tool import Tool
104-
from danswer.tools.tool import ToolResponse
107+
from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
108+
InternetSearchResponse,
109+
)
110+
from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
111+
InternetSearchTool,
112+
)
113+
from danswer.tools.tool_implementations.search.search_tool import (
114+
FINAL_CONTEXT_DOCUMENTS_ID,
115+
)
116+
from danswer.tools.tool_implementations.search.search_tool import (
117+
SEARCH_RESPONSE_SUMMARY_ID,
118+
)
119+
from danswer.tools.tool_implementations.search.search_tool import SearchResponseSummary
120+
from danswer.tools.tool_implementations.search.search_tool import SearchTool
121+
from danswer.tools.tool_implementations.search.search_tool import (
122+
SECTION_RELEVANCE_LIST_ID,
123+
)
105124
from danswer.tools.tool_runner import ToolCallFinalResult
106125
from danswer.tools.utils import compute_all_tool_tokens
107126
from danswer.tools.utils import explicit_tool_calling_supported
@@ -260,6 +279,7 @@ def _get_force_search_settings(
260279
| CustomToolResponse
261280
| MessageSpecificCitations
262281
| MessageResponseIDInfo
282+
| StreamStopInfo
263283
)
264284
ChatPacketStream = Iterator[ChatPacket]
265285

@@ -532,6 +552,13 @@ def stream_chat_message_objects(
532552
if not persona
533553
else PromptConfig.from_model(persona.prompts[0])
534554
)
555+
answer_style_config = AnswerStyleConfig(
556+
citation_config=CitationConfig(
557+
all_docs_useful=selected_db_search_docs is not None
558+
),
559+
document_pruning_config=document_pruning_config,
560+
structured_response_format=new_msg_req.structured_response_format,
561+
)
535562

536563
# find out what tools to use
537564
search_tool: SearchTool | None = None
@@ -550,13 +577,16 @@ def stream_chat_message_objects(
550577
llm=llm,
551578
fast_llm=fast_llm,
552579
pruning_config=document_pruning_config,
580+
answer_style_config=answer_style_config,
553581
selected_sections=selected_sections,
554582
chunks_above=new_msg_req.chunks_above,
555583
chunks_below=new_msg_req.chunks_below,
556584
full_doc=new_msg_req.full_doc,
557-
evaluation_type=LLMEvaluationType.BASIC
558-
if persona.llm_relevance_filter
559-
else LLMEvaluationType.SKIP,
585+
evaluation_type=(
586+
LLMEvaluationType.BASIC
587+
if persona.llm_relevance_filter
588+
else LLMEvaluationType.SKIP
589+
),
560590
)
561591
tool_dict[db_tool_model.id] = [search_tool]
562592
elif tool_cls.__name__ == ImageGenerationTool.__name__:
@@ -626,7 +656,11 @@ def stream_chat_message_objects(
626656
"Internet search tool requires a Bing API key, please contact your Danswer admin to get it added!"
627657
)
628658
tool_dict[db_tool_model.id] = [
629-
InternetSearchTool(api_key=bing_api_key)
659+
InternetSearchTool(
660+
api_key=bing_api_key,
661+
answer_style_config=answer_style_config,
662+
prompt_config=prompt_config,
663+
)
630664
]
631665

632666
continue
@@ -667,13 +701,7 @@ def stream_chat_message_objects(
667701
is_connected=is_connected,
668702
question=final_msg.message,
669703
latest_query_files=latest_query_files,
670-
answer_style_config=AnswerStyleConfig(
671-
citation_config=CitationConfig(
672-
all_docs_useful=selected_db_search_docs is not None
673-
),
674-
document_pruning_config=document_pruning_config,
675-
structured_response_format=new_msg_req.structured_response_format,
676-
),
704+
answer_style_config=answer_style_config,
677705
prompt_config=prompt_config,
678706
llm=(
679707
llm
@@ -777,7 +805,8 @@ def stream_chat_message_objects(
777805
response=custom_tool_response.tool_result,
778806
tool_name=custom_tool_response.tool_name,
779807
)
780-
808+
elif isinstance(packet, StreamStopInfo):
809+
pass
781810
else:
782811
if isinstance(packet, ToolCallFinalResult):
783812
tool_result = packet
@@ -807,6 +836,7 @@ def stream_chat_message_objects(
807836

808837
# Post-LLM answer processing
809838
try:
839+
logger.debug("Post-LLM answer processing")
810840
message_specific_citations: MessageSpecificCitations | None = None
811841
if reference_db_search_docs:
812842
message_specific_citations = _translate_citations(

0 commit comments

Comments
 (0)